{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Car Manual Explorer - GCP + MDB\n",
        "\n",
        "Using GCP (Document AI + Vertex AI) to parse, chunk and create embeddings from a car manual which are stored in an Atlas Collection. This will be utilised to build a Technical Document Intelligence System.\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "GE7lGFZvlFrP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 1. Create a GCP project and setup configs for it"
      ],
      "metadata": {
        "id": "GQiCVrhjl5Ms"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Installing GCP dependencies."
      ],
      "metadata": {
        "id": "NH5YadpElyuj"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zb_yXu07iiwG"
      },
      "outputs": [],
      "source": [
        "!pip install --upgrade --quiet google-cloud-storage==2.17.0\n",
        "!pip install --upgrade --quiet google-cloud-documentai==2.29.1\n",
        "!pip install --upgrade --quiet google-cloud-aiplatform==1.49.0\n",
        "!pip install --upgrade --quiet vertexai==1.49.0"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud auth application-default login\n",
        "!gcloud config set project <PROJECT_NAME>\n",
        "!gcloud config get-value project"
      ],
      "metadata": {
        "id": "wh_RAtqCl9Zu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!gcloud auth application-default set-quota-project <PROJECT_NAME>"
      ],
      "metadata": {
        "id": "6Cwp2p3vnBuu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#getting the Google Cloud Project ID and storing it in a variable\n",
        "PROJECT_ID_DETAILS = !gcloud config get-value project\n",
        "PROJECT_ID = PROJECT_ID_DETAILS[0]"
      ],
      "metadata": {
        "id": "-e0BWz7-n3ow"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#retrieve and store the Google Cloud Platform (GCP) project number\n",
        "PROJECT_NUMBER_DETAILS = !gcloud projects describe $PROJECT_ID --format=\"value(projectNumber)\"\n",
        "PROJECT_NUMBER = PROJECT_NUMBER_DETAILS[0]  # The project number is item 0 in the list returned by the gcloud command"
      ],
      "metadata": {
        "id": "YTpSc18B7Va5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2. Parse using Document AI"
      ],
      "metadata": {
        "id": "lc6db5suHxIL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**2.1** Utilizing **Document OCR** Processor from Document AI to parse the car manual."
      ],
      "metadata": {
        "id": "ZNtQHB3JOZuQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Splitting the Car Manual pdf into chunks of 15 pages to stay under Document AI's 30 page limit."
      ],
      "metadata": {
        "id": "EDkFB78aUEoR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "1. PDF SPLITTER FUNCTION"
      ],
      "metadata": {
        "id": "aS9Hpjh1wAkH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# 1. Ensure PyPDF2 is installed\n",
        "!pip install PyPDF2 tqdm\n",
        "\n",
        "# 2. Save the UPDATED PDF splitter code to a file\n",
        "with open('pdf_splitter.py', 'w') as f:\n",
        "    f.write('''\n",
        "from PyPDF2 import PdfReader, PdfWriter\n",
        "import os\n",
        "import math\n",
        "\n",
        "def split_pdf_for_document_ai(\n",
        "    input_pdf_path,\n",
        "    output_dir,\n",
        "    pages_per_chunk=15, # Reduced default for safety margin\n",
        "    start_page_number=1 # 1-based page number to start processing from\n",
        "    ):\n",
        "    \\\"\\\"\\\"\n",
        "    Split a PDF into smaller chunks for Document AI processing, optionally skipping initial pages.\n",
        "\n",
        "    Args:\n",
        "        input_pdf_path: Path to the input PDF file.\n",
        "        output_dir: Directory to save the split PDF files.\n",
        "        pages_per_chunk: Number of pages per chunk.\n",
        "        start_page_number: The 1-based page number to start splitting from.\n",
        "\n",
        "    Returns:\n",
        "        Tuple: (List of paths to the split PDF files, total_pages_processed)\n",
        "    \\\"\\\"\\\"\n",
        "    if not os.path.exists(output_dir):\n",
        "        os.makedirs(output_dir)\n",
        "\n",
        "    pdf_name = os.path.splitext(os.path.basename(input_pdf_path))[0]\n",
        "\n",
        "    try:\n",
        "        pdf = PdfReader(input_pdf_path)\n",
        "        total_source_pages = len(pdf.pages)\n",
        "    except Exception as e:\n",
        "        print(f\"Error reading PDF {input_pdf_path}: {e}\")\n",
        "        return [], 0\n",
        "\n",
        "    # Adjust for 0-based indexing and calculate the first page index to include\n",
        "    start_page_index = max(0, start_page_number - 1)\n",
        "\n",
        "    if start_page_index >= total_source_pages:\n",
        "        print(f\"Warning: Start page ({start_page_number}) is beyond the total number of pages ({total_source_pages}). No pages to process.\")\n",
        "        return [], 0\n",
        "\n",
        "    # Calculate the number of pages to actually process\n",
        "    pages_to_process_count = total_source_pages - start_page_index\n",
        "    num_chunks = math.ceil(pages_to_process_count / pages_per_chunk)\n",
        "\n",
        "    chunk_paths = []\n",
        "    total_pages_processed_in_chunks = 0\n",
        "\n",
        "    print(f\"Source PDF has {total_source_pages} pages. Processing from page {start_page_number} ({pages_to_process_count} pages).\")\n",
        "    print(f\"Splitting into {num_chunks} chunks of up to {pages_per_chunk} pages each.\")\n",
        "\n",
        "    for i in range(num_chunks):\n",
        "        output = PdfWriter()\n",
        "\n",
        "        # Calculate start and end page indices relative to the *source* PDF\n",
        "        chunk_start_index_rel_processed = i * pages_per_chunk\n",
        "        chunk_end_index_rel_processed = min((i + 1) * pages_per_chunk, pages_to_process_count)\n",
        "\n",
        "        # Map back to source page indices\n",
        "        source_start_page_index = start_page_index + chunk_start_index_rel_processed\n",
        "        source_end_page_index = start_page_index + chunk_end_index_rel_processed\n",
        "\n",
        "        pages_in_this_chunk = 0\n",
        "        for page_num_idx in range(source_start_page_index, source_end_page_index):\n",
        "            # Ensure we don't go beyond the actual number of pages\n",
        "            if page_num_idx < total_source_pages:\n",
        "                try:\n",
        "                    output.add_page(pdf.pages[page_num_idx])\n",
        "                    pages_in_this_chunk += 1\n",
        "                except Exception as e:\n",
        "                     print(f\"Error adding page {page_num_idx + 1} to chunk {i+1}: {e}\")\n",
        "\n",
        "\n",
        "        if pages_in_this_chunk > 0:\n",
        "            output_path = os.path.join(output_dir, f\"{pdf_name}_chunk_{i+1}_pages_{source_start_page_index+1}-{source_end_page_index}.pdf\")\n",
        "            try:\n",
        "                with open(output_path, \"wb\") as output_file:\n",
        "                    output.write(output_file)\n",
        "                chunk_paths.append(output_path)\n",
        "                total_pages_processed_in_chunks += pages_in_this_chunk\n",
        "                print(f\"Created chunk {i+1}/{num_chunks}: pages {source_start_page_index+1}-{source_end_page_index} -> {output_path}\")\n",
        "            except Exception as e:\n",
        "                print(f\"Error writing chunk {i+1} to {output_path}: {e}\")\n",
        "        else:\n",
        "            print(f\"Skipping chunk {i+1} as it contains no valid pages.\")\n",
        "\n",
        "\n",
        "    print(f\"Finished splitting. Total pages included in chunks: {total_pages_processed_in_chunks}\")\n",
        "    return chunk_paths, total_pages_processed_in_chunks\n",
        "''')"
      ],
      "metadata": {
        "id": "tVbH-vpFSYPx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Writing the function to send chunks of manual pdf to Document AI for **OCR Parsing**."
      ],
      "metadata": {
        "id": "MVbP1aXoUZNE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Keep parse_document function as is\n",
        "from google.api_core.client_options import ClientOptions\n",
        "from google.cloud import documentai_v1 as documentai\n",
        "from typing import List, Dict, Tuple, Any\n",
        "import tqdm\n",
        "import os # Make sure os is imported\n",
        "from PyPDF2 import PdfReader\n",
        "\n",
        "# Import the updated splitter\n",
        "from pdf_splitter import split_pdf_for_document_ai\n",
        "\n",
        "def parse_document(file_path, project_id, location, processor_id):\n",
        "    \"\"\"\n",
        "    Process a PDF document with Document AI. (Identical to original)\n",
        "    \"\"\"\n",
        "    from google.cloud import documentai # Ensure import is within function if needed\n",
        "\n",
        "    # Initialize Document AI client\n",
        "    # Consider initializing the client outside the function if called repeatedly\n",
        "    # for better performance, but this is safer for potential reuse.\n",
        "    client_options = ClientOptions(api_endpoint=f\"{location}-documentai.googleapis.com\")\n",
        "    client = documentai.DocumentProcessorServiceClient(client_options=client_options)\n",
        "\n",
        "\n",
        "    # Format the resource name\n",
        "    name = f\"projects/{project_id}/locations/{location}/processors/{processor_id}\"\n",
        "\n",
        "    # Read the file into memory\n",
        "    try:\n",
        "        with open(file_path, \"rb\") as file:\n",
        "            content = file.read()\n",
        "    except FileNotFoundError:\n",
        "        print(f\"Error: Chunk file not found at {file_path}\")\n",
        "        return None\n",
        "    except Exception as e:\n",
        "        print(f\"Error reading chunk file {file_path}: {e}\")\n",
        "        return None\n",
        "\n",
        "\n",
        "    # Configure the process request\n",
        "    document = documentai.RawDocument(content=content, mime_type=\"application/pdf\")\n",
        "    request = documentai.ProcessRequest(name=name, raw_document=document)\n",
        "\n",
        "    # Process the document\n",
        "    try:\n",
        "        result = client.process_document(request=request)\n",
        "        return result.document\n",
        "    except Exception as e:\n",
        "        print(f\"Error processing document {file_path} with Document AI: {e}\")\n",
        "        return None\n",
        "\n",
        "\n",
        "# --- REFACTORED Batch Processing Function ---\n",
        "def process_pdf_in_batches_ocr_only_list_output(\n",
        "    pdf_path: str,\n",
        "    project_id: str,\n",
        "    location: str,\n",
        "    ocr_processor_id: str,\n",
        "    output_dir: str = \"split_pdfs\",\n",
        "    pages_per_chunk: int = 15,\n",
        "    parallel_processing: bool = True,\n",
        "    max_concurrent: int = 3\n",
        ") -> Tuple[List[documentai.Document], int]:\n",
        "    \"\"\"\n",
        "    Process a PDF by splitting it and processing each chunk with Document AI OCR.\n",
        "\n",
        "    Args:\n",
        "        pdf_path: Path to the PDF file.\n",
        "        project_id: GCP project ID.\n",
        "        location: GCP location.\n",
        "        ocr_processor_id: Document AI OCR processor ID.\n",
        "        output_dir: Directory to save split PDFs.\n",
        "        pages_per_chunk: Number of pages per chunk.\n",
        "        parallel_processing: Whether to process chunks in parallel.\n",
        "        max_concurrent: Maximum number of concurrent API calls.\n",
        "\n",
        "    Returns:\n",
        "        Tuple: (List of processed Document AI document objects, expected_pages_processed)\n",
        "    \"\"\"\n",
        "    # Get the original number of pages for validation\n",
        "    try:\n",
        "        reader = PdfReader(pdf_path)\n",
        "        expected_pages_processed = len(reader.pages)\n",
        "    except Exception as e:\n",
        "        print(f\"Error reading original PDF for page count: {e}\")\n",
        "        expected_pages_processed = -1\n",
        "        return [], expected_pages_processed\n",
        "\n",
        "    # Split the PDF into chunks\n",
        "    print(f\"Splitting PDF '{pdf_path}' into chunks...\")\n",
        "    chunk_paths, total_pages_in_chunks = split_pdf_for_document_ai(\n",
        "        pdf_path, output_dir, pages_per_chunk\n",
        "    )\n",
        "\n",
        "    if not chunk_paths:\n",
        "        print(\"No PDF chunks were created. Exiting processing.\")\n",
        "        return [], expected_pages_processed\n",
        "\n",
        "    # Initialize list to store results\n",
        "    processed_documents = []\n",
        "    actual_pages_processed_api = 0\n",
        "\n",
        "    print(f\"\\nProcessing {len(chunk_paths)} PDF chunks with Document AI OCR...\")\n",
        "\n",
        "    if parallel_processing:\n",
        "        # Setup parallel processing of chunks\n",
        "        from concurrent.futures import ThreadPoolExecutor\n",
        "        import threading\n",
        "\n",
        "        # Thread-safe lock for printing\n",
        "        print_lock = threading.Lock()\n",
        "\n",
        "        # Define worker function with its own tqdm progress display\n",
        "        def process_chunk(idx_chunk_tuple):\n",
        "            idx, chunk_path = idx_chunk_tuple\n",
        "            with print_lock:\n",
        "                print(f\"\\nProcessing chunk {idx+1}/{len(chunk_paths)}: {os.path.basename(chunk_path)}\")\n",
        "\n",
        "            # Process with OCR processor\n",
        "            ocr_document = parse_document(\n",
        "                chunk_path, project_id, location, ocr_processor_id\n",
        "            )\n",
        "\n",
        "            result = {\"ocr_document\": None, \"pages_count\": 0, \"success\": False}\n",
        "\n",
        "            if ocr_document:\n",
        "                if hasattr(ocr_document, 'pages') and ocr_document.pages:\n",
        "                    pages_in_doc = len(ocr_document.pages)\n",
        "                    with print_lock:\n",
        "                        print(f\"  ✓ OCR processing successful for chunk {idx+1} ({pages_in_doc} pages)\")\n",
        "                    result = {\n",
        "                        \"ocr_document\": ocr_document,\n",
        "                        \"pages_count\": pages_in_doc,\n",
        "                        \"success\": True\n",
        "                    }\n",
        "                else:\n",
        "                    with print_lock:\n",
        "                        print(f\"  ✓ OCR processing successful for chunk {idx+1}, but no pages found.\")\n",
        "            else:\n",
        "                with print_lock:\n",
        "                    print(f\"  ✗ OCR processing failed for chunk {idx+1}\")\n",
        "\n",
        "            return result\n",
        "\n",
        "        # Process chunks in parallel with tqdm progress bar\n",
        "        with ThreadPoolExecutor(max_workers=max_concurrent) as executor:\n",
        "            results = list(tqdm.tqdm(\n",
        "                executor.map(process_chunk, enumerate(chunk_paths)),\n",
        "                total=len(chunk_paths),\n",
        "                desc=\"Processing Chunks\"\n",
        "            ))\n",
        "\n",
        "        # Collect results from parallel processing\n",
        "        for result in results:\n",
        "            if result[\"success\"]:\n",
        "                processed_documents.append(result[\"ocr_document\"])\n",
        "                actual_pages_processed_api += result[\"pages_count\"]\n",
        "    else:\n",
        "        # Original sequential processing with tqdm\n",
        "        for i, chunk_path in enumerate(tqdm.tqdm(chunk_paths, desc=\"Processing Chunks\")):\n",
        "            print(f\"\\nProcessing chunk {i+1}/{len(chunk_paths)}: {os.path.basename(chunk_path)}\")\n",
        "\n",
        "            # Process with OCR processor\n",
        "            ocr_document = parse_document(\n",
        "                chunk_path, project_id, location, ocr_processor_id\n",
        "            )\n",
        "\n",
        "            if ocr_document:\n",
        "                if hasattr(ocr_document, 'pages') and ocr_document.pages:\n",
        "                    pages_in_doc = len(ocr_document.pages)\n",
        "                    print(f\"  ✓ OCR processing successful ({pages_in_doc} pages found in result)\")\n",
        "                    processed_documents.append(ocr_document)\n",
        "                    actual_pages_processed_api += pages_in_doc\n",
        "                    # Debug first chunk\n",
        "                    if i == 0:\n",
        "                        print(\"\\n=== OCR Document Debug (First Processed Chunk) ===\")\n",
        "                        explore_ocr_document(ocr_document)\n",
        "                        print(\"===============================================\\n\")\n",
        "                else:\n",
        "                    print(f\"  ✓ OCR processing successful, but no pages found in the result.\")\n",
        "            else:\n",
        "                print(f\"  ✗ OCR processing failed for {os.path.basename(chunk_path)}\")\n",
        "\n",
        "    print(f\"\\nFinished processing chunks.\")\n",
        "    print(f\"Total pages expected to be processed: {expected_pages_processed}\")\n",
        "    print(f\"Total pages found in successful API results: {actual_pages_processed_api}\")\n",
        "\n",
        "    if expected_pages_processed > 0 and actual_pages_processed_api != expected_pages_processed:\n",
        "        print(f\"Warning: Mismatch between expected pages ({expected_pages_processed}) and pages in API results ({actual_pages_processed_api}). Check logs for errors.\")\n",
        "    elif expected_pages_processed == actual_pages_processed_api:\n",
        "        print(\"Page count validation successful.\")\n",
        "\n",
        "    return processed_documents, expected_pages_processed"
      ],
      "metadata": {
        "id": "WtmZF56tSyIT"
      },
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": [
        "Calling **OCR Parser**."
      ],
      "metadata": {
        "id": "SNnBcDwkUjrU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pdf_filename = \"/content/ABC Manual.pdf\" # Make sure this path is correct\n",
        "PROJECT_ID = \"ENTER_YOUR_ID_HERE\"\n",
        "LOCATION = \"us\"\n",
        "ocr_processor_id = \"ENTER_PROCESSOR_ID\"\n",
        "\n",
        "\n",
        "# Check if PDF exists before processing\n",
        "if not os.path.exists(pdf_filename):\n",
        "    print(f\"Error: PDF file not found at {pdf_filename}\")\n",
        "else:\n",
        "    # Process the PDF in batches, getting a list of documents\n",
        "    list_of_ocr_documents, expected_pages = process_pdf_in_batches_ocr_only_list_output(\n",
        "        pdf_path=pdf_filename,\n",
        "        project_id=PROJECT_ID,\n",
        "        location=LOCATION,\n",
        "        ocr_processor_id=ocr_processor_id,\n",
        "        output_dir=\"split_pdfs\",\n",
        "        pages_per_chunk=15\n",
        "    )\n",
        "\n",
        "    if list_of_ocr_documents:\n",
        "        print(f\"\\nSuccessfully processed {len(list_of_ocr_documents)} chunks.\")\n",
        "        # You can now pass list_of_ocr_documents to the updated text anchor processor\n",
        "    else:\n",
        "        print(\"\\nNo documents were successfully processed.\")"
      ],
      "metadata": {
        "id": "1r5MpikzWAZu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "Inspect **OCR Document**.\n"
      ],
      "metadata": {
        "id": "HNZLKZu9xaey"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def explore_ocr_document(ocr_document):\n",
        "    \"\"\"Thoroughly explore OCR document structure to find where text is stored\"\"\"\n",
        "    print(\"\\n===== EXPLORING OCR DOCUMENT STRUCTURE =====\")\n",
        "\n",
        "    # Check document-level text\n",
        "    if hasattr(ocr_document, 'text'):\n",
        "        print(f\"Document has text attribute with {len(ocr_document.text)} characters\")\n",
        "        print(f\"Example: {ocr_document.text[:100]}...\")\n",
        "    else:\n",
        "        print(\"Document does NOT have a text attribute\")\n",
        "\n",
        "    # Check document attributes\n",
        "    print(f\"\\nDocument attributes: {[attr for attr in dir(ocr_document) if not attr.startswith('_') and not callable(getattr(ocr_document, attr))]}\")\n",
        "\n",
        "    # Check first page in detail\n",
        "    if hasattr(ocr_document, 'pages') and len(ocr_document.pages) > 0:\n",
        "        first_page = ocr_document.pages[0]\n",
        "        print(f\"\\nFirst page attributes: {[attr for attr in dir(first_page) if not attr.startswith('_') and not callable(getattr(first_page, attr))]}\")\n",
        "\n",
        "        # Look for blocks\n",
        "        if hasattr(first_page, 'blocks') and len(first_page.blocks) > 0:\n",
        "            print(f\"\\nPage has {len(first_page.blocks)} blocks\")\n",
        "            first_block = first_page.blocks[0]\n",
        "            print(f\"First block attributes: {[attr for attr in dir(first_block) if not attr.startswith('_') and not callable(getattr(first_block, attr))]}\")\n",
        "\n",
        "            # Try to find text in blocks\n",
        "            if hasattr(first_block, 'text'):\n",
        "                print(f\"Block has text: {first_block.text[:50]}...\")\n",
        "\n",
        "            # Try to find text anchor\n",
        "            if hasattr(first_block, 'layout') and hasattr(first_block.layout, 'text_anchor'):\n",
        "                print(\"Block has text_anchor - text might be referenced via offsets\")\n",
        "\n",
        "                text_anchor = first_block.layout.text_anchor\n",
        "                print(f\"Text anchor attributes: {[attr for attr in dir(text_anchor) if not attr.startswith('_') and not callable(getattr(text_anchor, attr))]}\")\n",
        "\n",
        "                if hasattr(text_anchor, 'text_segments') and len(text_anchor.text_segments) > 0:\n",
        "                    segment = text_anchor.text_segments[0]\n",
        "                    print(f\"Text segment attributes: {[attr for attr in dir(segment) if not attr.startswith('_') and not callable(getattr(segment, attr))]}\")\n",
        "\n",
        "                    # Try to extract text from segment\n",
        "                    if hasattr(segment, 'start_index') and hasattr(segment, 'end_index') and hasattr(ocr_document, 'text'):\n",
        "                        start = segment.start_index\n",
        "                        end = segment.end_index\n",
        "                        print(f\"Text segment points to: {ocr_document.text[start:end]}\")\n",
        "\n",
        "            # Look for paragraphs\n",
        "            if hasattr(first_block, 'paragraphs') and len(first_block.paragraphs) > 0:\n",
        "                print(f\"\\nBlock has {len(first_block.paragraphs)} paragraphs\")\n",
        "                first_para = first_block.paragraphs[0]\n",
        "                print(f\"First paragraph attributes: {[attr for attr in dir(first_para) if not attr.startswith('_') and not callable(getattr(first_para, attr))]}\")\n",
        "\n",
        "                if hasattr(first_para, 'text'):\n",
        "                    print(f\"Paragraph has text: {first_para.text[:50]}...\")\n",
        "\n",
        "                # Check for words in paragraph\n",
        "                if hasattr(first_para, 'words') and len(first_para.words) > 0:\n",
        "                    print(f\"Paragraph has {len(first_para.words)} words\")\n",
        "                    first_word = first_para.words[0]\n",
        "                    print(f\"First word attributes: {[attr for attr in dir(first_word) if not attr.startswith('_') and not callable(getattr(first_word, attr))]}\")\n",
        "\n",
        "                    if hasattr(first_word, 'text'):\n",
        "                        print(f\"Word has text: {first_word.text}\")\n",
        "\n",
        "    # Check text_segments if available directly\n",
        "    if hasattr(ocr_document, 'text_segments'):\n",
        "        print(f\"\\nDocument has {len(ocr_document.text_segments)} text segments\")\n",
        "        if len(ocr_document.text_segments) > 0:\n",
        "            print(f\"First segment: {ocr_document.text_segments[0]}\")\n",
        "\n",
        "    print(\"\\n=== RECOMMENDATION FOR TEXT EXTRACTION ===\")\n",
        "    # Recommend the best approach based on what was found\n",
        "    if hasattr(ocr_document, 'text') and ocr_document.text:\n",
        "        # If the document has full text and blocks have text_anchors, use that approach\n",
        "        print(\"1. RECOMMENDED: Extract text using text_anchors that reference document.text\")\n",
        "    elif hasattr(first_page, 'blocks') and hasattr(first_block, 'paragraphs'):\n",
        "        if hasattr(first_para, 'text'):\n",
        "            print(\"2. RECOMMENDED: Extract text from block -> paragraph -> text\")\n",
        "        elif hasattr(first_para, 'words'):\n",
        "            print(\"3. RECOMMENDED: Extract text from block -> paragraph -> words -> text\")\n",
        "    else:\n",
        "        print(\"4. FALLBACK: Use simpler text extraction approach\")\n",
        "\n",
        "    print(\"==========================================\\n\")\n",
        "\n",
        "    return True"
      ],
      "metadata": {
        "id": "ogsD1OlD_vVO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "VALIDATION FUNCTION"
      ],
      "metadata": {
        "id": "4DzrAp4GyoBO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from PyPDF2 import PdfReader\n",
        "import difflib\n",
        "\n",
        "def validate_parsed_text(\n",
        "    combined_data: Dict[Tuple[int, float], Dict[str, Any]],\n",
        "    original_pdf_path: str,\n",
        "    start_page_number: int,\n",
        "    expected_pages_processed: int,\n",
        "    num_sample_pages: int = 3\n",
        "    ):\n",
        "    \"\"\"\n",
        "    Validates the parsed text against the original PDF and checks structure.\n",
        "\n",
        "    Args:\n",
        "        combined_data: The output from process_ocr_list_with_text_anchors.\n",
        "        original_pdf_path: Path to the source PDF.\n",
        "        start_page_number: 1-based page number where processing started.\n",
        "        expected_pages_processed: Number of pages expected to be processed.\n",
        "        num_sample_pages: How many pages to sample for text comparison.\n",
        "    \"\"\"\n",
        "    print(\"\\n--- Starting Validation ---\")\n",
        "\n",
        "    if not combined_data:\n",
        "        print(\"Validation Error: No combined data provided.\")\n",
        "        return\n",
        "\n",
        "    processed_pages = sorted(list(set(key[0] for key in combined_data.keys())))\n",
        "    if not processed_pages:\n",
        "        print(\"Validation Error: No pages found in combined data.\")\n",
        "        return\n",
        "\n",
        "    actual_pages_count = len(processed_pages)\n",
        "    min_page = min(processed_pages)\n",
        "    max_page = max(processed_pages)\n",
        "\n",
        "    print(f\"Found {actual_pages_count} unique page numbers in processed data (Range: {min_page}-{max_page}).\")\n",
        "    if expected_pages_processed > 0 :\n",
        "         if actual_pages_count == expected_pages_processed:\n",
        "             print(f\"Page Count Check: OK (Matches expected {expected_pages_processed} pages)\")\n",
        "         else:\n",
        "              print(f\"Page Count Check: MISMATCH (Expected {expected_pages_processed}, Found {actual_pages_count})\")\n",
        "    else:\n",
        "         print(\"Page Count Check: Skipped (Could not get expected count from source PDF)\")\n",
        "\n",
        "\n",
        "    # 2. Text Comparison for Sample Pages\n",
        "    print(f\"\\nComparing text for {num_sample_pages} sample pages...\")\n",
        "    try:\n",
        "        pdf = PdfReader(original_pdf_path)\n",
        "        total_source_pages = len(pdf.pages)\n",
        "\n",
        "        # Select sample page numbers from the processed range\n",
        "        sample_page_indices = []\n",
        "        if actual_pages_count > 0:\n",
        "            step = max(1, actual_pages_count // num_sample_pages)\n",
        "            sample_page_indices = [processed_pages[i] for i in range(0, actual_pages_count, step)]\n",
        "            # Ensure first and last processed pages are included if possible\n",
        "            if min_page not in sample_page_indices:\n",
        "                 sample_page_indices.insert(0, min_page)\n",
        "            if max_page not in sample_page_indices and max_page != min_page:\n",
        "                  sample_page_indices.append(max_page)\n",
        "            sample_page_indices = sorted(list(set(sample_page_indices)))[:num_sample_pages] # Ensure unique and limit count\n",
        "\n",
        "\n",
        "        for page_num_to_check in sample_page_indices:\n",
        "            print(f\"\\n--- Comparing Page {page_num_to_check} ---\")\n",
        "            page_num_0_based = page_num_to_check - 1\n",
        "\n",
        "            # Extract text from original PDF using PyPDF2\n",
        "            original_text = \"\"\n",
        "            if 0 <= page_num_0_based < total_source_pages:\n",
        "                try:\n",
        "                    page = pdf.pages[page_num_0_based]\n",
        "                    original_text = page.extract_text()\n",
        "                    if not original_text:\n",
        "                         original_text = \"[PyPDF2 extracted no text]\"\n",
        "                    else:\n",
        "                         original_text = re.sub(r'\\s+', ' ', original_text).strip() # Normalize whitespace\n",
        "                except Exception as e:\n",
        "                    original_text = f\"[Error extracting text with PyPDF2: {e}]\"\n",
        "            else:\n",
        "                 original_text = \"[Page number out of range for original PDF]\"\n",
        "\n",
        "\n",
        "            # Extract text from processed data\n",
        "            processed_elements_on_page = sorted(\n",
        "                [elem for key, elem in combined_data.items() if key[0] == page_num_to_check],\n",
        "                key=lambda x: x['id'] # Sort by element ID for consistent order\n",
        "            )\n",
        "            processed_text = \" \".join(elem['text'] for elem in processed_elements_on_page)\n",
        "            processed_text = re.sub(r'\\s+', ' ', processed_text).strip() # Normalize\n",
        "\n",
        "\n",
        "            # Compare texts\n",
        "            print(f\"Original PDF (PyPDF2):\\n\\\"{original_text[:300]}...\\\"\")\n",
        "            print(f\"\\nProcessed (Document AI):\\n\\\"{processed_text[:300]}...\\\"\")\n",
        "\n",
        "            # Simple similarity check (ratio)\n",
        "            similarity = difflib.SequenceMatcher(None, original_text.lower(), processed_text.lower()).ratio()\n",
        "            print(f\"\\nText Similarity Ratio: {similarity:.2f}\")\n",
        "            if similarity < 0.7: # Threshold can be adjusted\n",
        "                print(\"WARNING: Low similarity detected. Manual review recommended.\")\n",
        "            else:\n",
        "                print(\"Similarity Check: OK\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"Error during text comparison: {e}\")\n",
        "\n",
        "    # 3. Check for empty text elements\n",
        "    empty_text_elements = [elem['id'] for elem in combined_data.values() if not elem['text'].strip()]\n",
        "    if empty_text_elements:\n",
        "        print(f\"\\nWarning: Found {len(empty_text_elements)} elements with empty text (e.g., {empty_text_elements[:5]}).\")\n",
        "    else:\n",
        "        print(\"\\nEmpty Text Check: OK (No empty text elements found)\")\n",
        "\n",
        "    # 4. Check section hierarchy consistency (basic check)\n",
        "    print(\"\\nChecking Section Hierarchy (Sample):\")\n",
        "    hierarchy_issues = 0\n",
        "    last_hierarchy = {}\n",
        "    sample_elements = list(combined_data.values())[::max(1, len(combined_data)//20)] # Sample ~5% of elements\n",
        "\n",
        "    for elem in sample_elements:\n",
        "         current_hierarchy = elem.get('section_hierarchy', {})\n",
        "         # Simple check: lower level heading should not appear without a higher level one set previously or concurrently\n",
        "         if current_hierarchy.get(3) and not current_hierarchy.get(2) and not last_hierarchy.get(2):\n",
        "              hierarchy_issues += 1\n",
        "              # print(f\"  Potential hierarchy issue at element {elem['id']} (page {elem['page_number']}): Subheading found without recent Heading.\")\n",
        "         if current_hierarchy.get(2) and not current_hierarchy.get(1) and not last_hierarchy.get(1):\n",
        "              hierarchy_issues += 1\n",
        "              # print(f\"  Potential hierarchy issue at element {elem['id']} (page {elem['page_number']}): Heading found without recent Title.\")\n",
        "         last_hierarchy = current_hierarchy # Update for next comparison\n",
        "\n",
        "\n",
        "    if hierarchy_issues > 0:\n",
        "         print(f\"WARNING: Found {hierarchy_issues} potential hierarchy inconsistencies in sample. Review 'section_hierarchy' in data.\")\n",
        "    else:\n",
        "         print(\"Section Hierarchy Check: OK (Basic consistency in sample)\")\n",
        "\n",
        "    print(\"\\n--- Validation Complete ---\")"
      ],
      "metadata": {
        "id": "DoNUR5thypaM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "TEXT ANCHOR\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "j30tGHkwPXrP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import re\n",
        "from typing import List, Dict, Tuple, Any\n",
        "from google.cloud import documentai_v1 as documentai # Ensure import if needed\n",
        "\n",
        "# --- REFACTORED Text Anchor Processing Function ---\n",
        "def process_ocr_list_with_text_anchors(\n",
        "    ocr_documents: List[documentai.Document]\n",
        ") -> Dict[Tuple[int, float], Dict[str, Any]]:\n",
        "    \"\"\"\n",
        "    Extract text and infer structure from a list of OCR document results using text anchors.\n",
        "    Enhanced for car manual structure.\n",
        "\n",
        "    Args:\n",
        "        ocr_documents: List of Document AI results, one per processed chunk.\n",
        "\n",
        "    Returns:\n",
        "        Dictionary of position tuples to element dictionaries for the entire document.\n",
        "    \"\"\"\n",
        "    print(f\"Processing {len(ocr_documents)} OCR document chunks using text anchors...\")\n",
        "    combined_data = {}\n",
        "    element_id_counter = 0\n",
        "    current_doc_page_offset = 0\n",
        "\n",
        "    # Track document hierarchical structure ACROSS chunks\n",
        "    current_title = None\n",
        "    current_section = None\n",
        "    current_heading = None\n",
        "    current_subheading = None\n",
        "\n",
        "    # Car manual specific patterns\n",
        "    car_manual_section_patterns = [\n",
        "        r'(?i)^(chapter|section)\\s+\\d+',\n",
        "        r'(?i)^(introduction|overview|specifications|maintenance|troubleshooting|appendix)'\n",
        "    ]\n",
        "\n",
        "    warning_patterns = [\n",
        "        r'(?i)^(warning|caution|note|danger|attention|important):',\n",
        "        r'(?i)^(warning|caution|note|danger|attention|important)$'\n",
        "    ]\n",
        "\n",
        "    procedure_patterns = [\n",
        "        r'(?i)^\\d+\\.\\s+[A-Z]',  # Numbered steps starting with capital letter\n",
        "        r'(?i)^to\\s+[a-z]+',    # \"To do something\" procedure titles\n",
        "    ]\n",
        "\n",
        "    spec_patterns = [\n",
        "        r'(?i)specifications',\n",
        "        r'(?i)technical data',\n",
        "        r'[\\d.]+ (mm|cm|in|inches|kg|lbs|L|liters|gal|gallons)'  # Measurements\n",
        "    ]\n",
        "\n",
        "    # Ignore patterns for context (avoid these as section headers)\n",
        "    ignore_patterns = [\n",
        "        r'Part Number:', r'Copyright', r'Ford Motor Company', r'www\\.ford\\.com',\n",
        "        r'\\d{8}', r'^\\d{6}$', r'@', r'^\\s*$'\n",
        "    ]\n",
        "\n",
        "    # Process each document (chunk) in the list\n",
        "    for doc_idx, ocr_document in enumerate(ocr_documents):\n",
        "        # Check for document text\n",
        "        if not hasattr(ocr_document, 'text') or not ocr_document.text:\n",
        "            print(f\"Warning: Document chunk {doc_idx+1} has no text attribute. Skipping.\")\n",
        "            continue\n",
        "\n",
        "        full_text = ocr_document.text # Text is relative to THIS document chunk\n",
        "        print(f\"Processing document chunk {doc_idx+1} with {len(full_text)} characters.\")\n",
        "\n",
        "        # Check if OCR document has pages\n",
        "        if not hasattr(ocr_document, 'pages') or not ocr_document.pages:\n",
        "            print(f\"Warning: Document chunk {doc_idx+1} has no pages. Skipping.\")\n",
        "            continue\n",
        "\n",
        "        pages_in_current_doc = len(ocr_document.pages)\n",
        "\n",
        "        # Process each page within the current document chunk\n",
        "        for page_idx_in_doc, page in enumerate(ocr_document.pages):\n",
        "            # Calculate the absolute page number based on the progress\n",
        "            absolute_page_num = current_doc_page_offset + page_idx_in_doc + 1\n",
        "\n",
        "            # Get blocks on the page\n",
        "            if not hasattr(page, 'blocks') or not page.blocks:\n",
        "                continue\n",
        "\n",
        "            # Process each block on the page\n",
        "            for block_idx, block in enumerate(page.blocks):\n",
        "                if not hasattr(block, 'layout') or not hasattr(block.layout, 'text_anchor'):\n",
        "                    continue\n",
        "\n",
        "                text_anchor = block.layout.text_anchor\n",
        "                if not hasattr(text_anchor, 'text_segments') or not text_anchor.text_segments:\n",
        "                    continue\n",
        "\n",
        "                # Get text segments within the block\n",
        "                block_text_parts = []\n",
        "                segment_start_indices = []\n",
        "\n",
        "                for segment in text_anchor.text_segments:\n",
        "                    if not hasattr(segment, 'start_index') or not hasattr(segment, 'end_index'):\n",
        "                        continue\n",
        "\n",
        "                    start_index = segment.start_index\n",
        "                    end_index = segment.end_index\n",
        "\n",
        "                    if start_index < 0 or end_index > len(full_text) or start_index >= end_index:\n",
        "                        continue\n",
        "\n",
        "                    # Extract text using indices relative to the CURRENT document's text\n",
        "                    segment_text = full_text[start_index:end_index]\n",
        "                    block_text_parts.append(segment_text)\n",
        "                    segment_start_indices.append(start_index)\n",
        "\n",
        "                # Combine text parts for the block and clean it\n",
        "                text = \"\".join(block_text_parts).strip()\n",
        "                text = re.sub(r'\\s+', ' ', text) # Normalize whitespace\n",
        "\n",
        "                if not text: # Skip empty blocks\n",
        "                    continue\n",
        "\n",
        "                # Use the start index of the *first* segment for positioning/sorting\n",
        "                sort_pos = (segment_start_indices[0] / 1000000.0) if segment_start_indices else (block_idx / 1000.0)\n",
        "                position_key = (absolute_page_num, sort_pos)\n",
        "\n",
        "                # Enhanced element type detection for car manuals\n",
        "                element_type = \"PARAGRAPH\"\n",
        "                is_potential_header = False\n",
        "\n",
        "                # Car manual specific classifications\n",
        "                if any(re.search(pattern, text) for pattern in warning_patterns):\n",
        "                    element_type = \"WARNING\"\n",
        "                elif any(re.search(pattern, text) for pattern in spec_patterns):\n",
        "                    element_type = \"SPECIFICATION\"\n",
        "                elif any(re.search(pattern, text) for pattern in procedure_patterns):\n",
        "                    element_type = \"PROCEDURE\"\n",
        "                # Hierarchy patterns\n",
        "                elif text.isupper() and len(text.split()) <= 6 and len(text) > 3:\n",
        "                    if any(re.search(pattern, text) for pattern in car_manual_section_patterns):\n",
        "                        element_type = \"SECTION\"\n",
        "                        is_potential_header = True\n",
        "                    elif not any(re.search(pattern, text, re.IGNORECASE) for pattern in ignore_patterns):\n",
        "                        element_type = \"HEADING\"\n",
        "                        is_potential_header = True\n",
        "                elif (re.match(r'^\\d+\\.', text.strip()) or\n",
        "                      (re.match(r'^[A-Z][a-zA-Z\\s]+:', text.strip()) and len(text) < 100)) and \\\n",
        "                     not any(re.search(pattern, text, re.IGNORECASE) for pattern in ignore_patterns):\n",
        "                    element_type = \"HEADING\"\n",
        "                    is_potential_header = True\n",
        "                elif re.match(r'^[A-Z].{5,50}:$', text.strip()) and \\\n",
        "                     not any(re.search(pattern, text, re.IGNORECASE) for pattern in ignore_patterns):\n",
        "                    element_type = \"SUBHEADING\"\n",
        "                    is_potential_header = True\n",
        "                elif text.strip().startswith(('•', '-', '*')) or re.match(r'^\\d+[\\.)]', text.strip()):\n",
        "                    element_type = \"LIST_ITEM\"\n",
        "                elif bool(re.search(r'figure|fig\\.', text, re.IGNORECASE)):\n",
        "                    element_type = \"FIGURE\"\n",
        "                elif text.count('|') > 2 or text.count('\\t') > 2 or (':' in text and text.count('.') > 3):\n",
        "                    element_type = \"TABLE\"\n",
        "\n",
        "                # Update hierarchy if this element is a header type\n",
        "                if element_type == 'SECTION':\n",
        "                    current_section = text\n",
        "                    current_heading = None  # Reset lower levels\n",
        "                    current_subheading = None\n",
        "                elif element_type == 'HEADING':\n",
        "                    current_heading = text\n",
        "                    current_subheading = None  # Reset lower level\n",
        "                elif element_type == 'SUBHEADING':\n",
        "                    current_subheading = text\n",
        "\n",
        "                # Construct current hierarchy snapshot\n",
        "                current_hierarchy_snapshot = {}\n",
        "                if current_section and not any(re.search(pattern, current_section, re.IGNORECASE) for pattern in ignore_patterns):\n",
        "                    current_hierarchy_snapshot[1] = current_section\n",
        "                if current_heading and not any(re.search(pattern, current_heading, re.IGNORECASE) for pattern in ignore_patterns):\n",
        "                    current_hierarchy_snapshot[2] = current_heading\n",
        "                if current_subheading and not any(re.search(pattern, current_subheading, re.IGNORECASE) for pattern in ignore_patterns):\n",
        "                    current_hierarchy_snapshot[3] = current_subheading\n",
        "\n",
        "                # Bounding box\n",
        "                bbox = {\"x1\": 0.1, \"y1\": block_idx / 100.0, \"x2\": 0.9, \"y2\": (block_idx / 100.0) + 0.05}\n",
        "                if hasattr(block.layout, 'bounding_poly') and block.layout.bounding_poly.normalized_vertices:\n",
        "                    vertices = block.layout.bounding_poly.normalized_vertices\n",
        "                    if len(vertices) >= 4:\n",
        "                        bbox = {\n",
        "                            \"x1\": vertices[0].x,\n",
        "                            \"y1\": vertices[0].y,\n",
        "                            \"x2\": vertices[2].x,\n",
        "                            \"y2\": vertices[2].y\n",
        "                        }\n",
        "\n",
        "                # Create the structured element\n",
        "                element = {\n",
        "                    \"bbox\": bbox,\n",
        "                    \"type\": element_type,\n",
        "                    \"confidence\": block.layout.confidence if hasattr(block.layout, 'confidence') else 0.9,\n",
        "                    \"page_number\": absolute_page_num,\n",
        "                    \"text\": text,\n",
        "                    \"id\": f\"element_{element_id_counter}\",\n",
        "                    \"section_hierarchy\": current_hierarchy_snapshot,\n",
        "                    \"doc_chunk_index\": doc_idx\n",
        "                }\n",
        "                element_id_counter += 1\n",
        "\n",
        "                # Store in combined data\n",
        "                if position_key in combined_data:\n",
        "                    print(f\"Warning: Position key collision for page {absolute_page_num}. Overwriting.\")\n",
        "                combined_data[position_key] = element\n",
        "\n",
        "        # Update the page offset for the next document chunk\n",
        "        current_doc_page_offset += pages_in_current_doc\n",
        "\n",
        "    print(f\"Finished processing text anchors. Extracted {len(combined_data)} structured elements.\")\n",
        "    return combined_data"
      ],
      "metadata": {
        "id": "SYoYSsFPpMuZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "EXECUTE VALIDATION\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "S7EpMmgQEmpZ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "START_PAGE = 1\n",
        "\n",
        "if list_of_ocr_documents:\n",
        "    # Process the list of documents to extract structured text\n",
        "    combined_data_structured = process_ocr_list_with_text_anchors(\n",
        "        ocr_documents=list_of_ocr_documents\n",
        "    )\n",
        "\n",
        "    # Validate the results\n",
        "    if combined_data_structured:\n",
        "        validate_parsed_text(\n",
        "            combined_data=combined_data_structured,\n",
        "            original_pdf_path=pdf_filename,\n",
        "            start_page_number=START_PAGE,\n",
        "            expected_pages_processed=expected_pages # Pass the expected count from the batch processor\n",
        "        )\n",
        "\n",
        "        # You can now proceed to the chunking steps using combined_data_structured\n",
        "        # For example:\n",
        "        # initial_chunks = create_meaningful_chunks_ocr(combined_data_structured, ...)\n",
        "        print(\"\\nReady to proceed to chunking using 'combined_data_structured'.\")\n",
        "\n",
        "    else:\n",
        "        print(\"\\nText anchor processing failed to produce data. Cannot validate.\")\n",
        "else:\n",
        "    print(\"\\nSkipping text anchor processing and validation as OCR processing yielded no documents.\")"
      ],
      "metadata": {
        "id": "b2a7J-d6Epm6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 3. Chunking"
      ],
      "metadata": {
        "id": "9g2-C6EGzNx3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install --quiet nltk"
      ],
      "metadata": {
        "id": "4fWdKDfT98KB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!python -m nltk.downloader punkt"
      ],
      "metadata": {
        "id": "4pQ37-zS99ZH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import re\n",
        "import copy\n",
        "import datetime\n",
        "import json\n",
        "import gc\n",
        "import os\n",
        "from typing import List, Dict, Tuple, Any, Generator, Set, Optional\n",
        "from tqdm import tqdm\n",
        "\n",
        "# --- Optional: NLTK for better sentence splitting ---\n",
        "USE_NLTK = True # Set to True if you have nltk installed\n",
        "try:\n",
        "    import nltk\n",
        "    # Uncomment the next line if you have downloaded the 'punkt' resource\n",
        "    nltk.download('punkt', quiet=True)\n",
        "    from nltk.tokenize import sent_tokenize\n",
        "    if not USE_NLTK:\n",
        "        print(\"NLTK found, but USE_NLTK is False. Using regex sentence splitting.\")\n",
        "    else:\n",
        "        print(\"Using NLTK for sentence splitting.\")\n",
        "except ImportError:\n",
        "    print(\"NLTK not found or 'punkt' not downloaded. Using regex-based sentence splitting.\")\n",
        "    USE_NLTK = False\n",
        "    def sent_tokenize(text): # Basic regex fallback\n",
        "        sentences = re.split(r'(?<=[.!?])\\s+', text)\n",
        "        return [s for s in sentences if s] # Filter empty strings\n",
        "\n",
        "# --- Configuration ---\n",
        "MAX_CHUNK_SIZE = 1800  # Target maximum characters\n",
        "MIN_CHUNK_SIZE = 500   # Target minimum characters (Increased for better context)\n",
        "OVERLAP_SENTENCE_COUNT = 1 # Number of sentences for overlap\n",
        "HEADING_SPLIT_THRESHOLD_FACTOR = 0.8 # Only split before headings if chunk size > 80% of max_chunk_size\n",
        "AGGRESSIVE_MERGE_THRESHOLD = 150     # Chunks smaller than this will be merged more aggressively\n",
        "\n",
        "\n",
        "\n",
        "# --- Memory-Optimized Chunking Implementation ---\n",
        "# (process_elements_in_batches remains largely the same, ensuring context is passed correctly)\n",
        "def process_elements_in_batches(\n",
        "    combined_data: Dict[Tuple[int, float], Dict[str, Any]],\n",
        "    max_chunk_size: int = MAX_CHUNK_SIZE,\n",
        "    min_chunk_size: int = MIN_CHUNK_SIZE,\n",
        "    overlap_sentences: int = OVERLAP_SENTENCE_COUNT,\n",
        "    batch_size: int = 1000\n",
        ") -> List[Dict[str, Any]]:\n",
        "    \"\"\"\n",
        "    Memory-optimized function to process large manuals in batches.\n",
        "    (Context passing updated to remove procedure state)\n",
        "    \"\"\"\n",
        "    if not combined_data:\n",
        "        return []\n",
        "\n",
        "    sorted_keys = sorted(combined_data.keys())\n",
        "    total_elements = len(sorted_keys)\n",
        "\n",
        "    print(f\"Processing {total_elements} elements in batches of {batch_size}...\")\n",
        "\n",
        "    all_chunks = []\n",
        "    last_chunk_last_sentences = [] # Store sentences for overlap\n",
        "    last_elements_context = {} # Preserve hierarchy context\n",
        "\n",
        "    page_ranges = {key[0] for key in sorted_keys}\n",
        "    total_pages = len(page_ranges)\n",
        "    print(f\"Manual spans {total_pages} pages across {total_elements} elements.\")\n",
        "\n",
        "    for i in range(0, total_elements, batch_size):\n",
        "        batch_end = min(i + batch_size, total_elements)\n",
        "        batch_keys = sorted_keys[i:batch_end]\n",
        "        batch_elements = {key: combined_data[key] for key in batch_keys}\n",
        "\n",
        "        batch_pages = set(key[0] for key in batch_keys)\n",
        "        min_page, max_page = min(batch_pages), max(batch_pages)\n",
        "\n",
        "        print(f\"\\nProcessing batch {i//batch_size + 1}/{(total_elements+batch_size-1)//batch_size}: \"\n",
        "              f\"Elements {i+1}-{batch_end} (Pages {min_page}-{max_page})\")\n",
        "\n",
        "        batch_chunks = create_enriched_chunks(\n",
        "            batch_elements,\n",
        "            max_chunk_size,\n",
        "            min_chunk_size,\n",
        "            overlap_sentences,\n",
        "            context_from_previous_batch=last_elements_context,\n",
        "            overlap_sentences_from_previous=last_chunk_last_sentences\n",
        "        )\n",
        "\n",
        "        # Save the last chunk's ending sentences for potential overlap\n",
        "        if batch_chunks:\n",
        "            last_chunk_text = batch_chunks[-1]['text']\n",
        "            sentences = sent_tokenize(last_chunk_text)\n",
        "            last_chunk_last_sentences = sentences[-overlap_sentences:] if sentences else []\n",
        "\n",
        "            # Save hierarchy context for continuation in next batch\n",
        "            context_levels = ['heading_level_1', 'heading_level_2', 'heading_level_3']\n",
        "            last_elements_context = {level: batch_chunks[-1].get(level) for level in context_levels}\n",
        "\n",
        "        all_chunks.extend(batch_chunks)\n",
        "\n",
        "        batch_elements = None # Free memory\n",
        "        gc.collect()\n",
        "\n",
        "        print(f\"Batch complete. Generated {len(batch_chunks)} chunks. \"\n",
        "              f\"Total chunks so far: {len(all_chunks)}\")\n",
        "\n",
        "    print(\"\\nFinalizing all chunks...\")\n",
        "    final_chunks = post_process_chunks(all_chunks)\n",
        "\n",
        "    print(f\"Processing complete. Total chunks: {len(final_chunks)}\")\n",
        "    return final_chunks"
      ],
      "metadata": {
        "id": "fx4KAvni-LXx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "MERGE CHUNK - HELPER FUNCTION\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "lSC54VFQKBZl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# --- Helper Function for Merging Chunks ---\n",
        "def _perform_merge(prev_chunk: Dict[str, Any], current_chunk: Dict[str, Any], overlap_sentences: int):\n",
        "    \"\"\"\n",
        "    Merges the current_chunk into the prev_chunk in place.\n",
        "    Updates text, page numbers, features, and metadata.\n",
        "    Returns the last sentences of the merged chunk for potential overlap.\n",
        "    \"\"\"\n",
        "    # Merge text with a clear separator\n",
        "    prev_chunk['text'] += \"\\n\\n---\\n\\n\" + current_chunk['text'] # Add a distinct separator\n",
        "\n",
        "    # Update page numbers\n",
        "    prev_chunk['page_numbers'] = sorted(list(set(prev_chunk['page_numbers'] + current_chunk['page_numbers'])))\n",
        "\n",
        "    # Re-run feature extraction and metadata update on the *merged* content\n",
        "    # Note: This assumes extract_chunk_features works purely on text. If it relied\n",
        "    #       on element details, merging features would be more complex.\n",
        "    updated_features = extract_chunk_features(prev_chunk['text'])\n",
        "    prev_chunk.update(updated_features) # Overwrite/add new features\n",
        "\n",
        "    # Combine content types carefully\n",
        "    prev_chunk['content_type'] = sorted(list(set(prev_chunk.get('content_type', []) + current_chunk.get('content_type', []))))\n",
        "\n",
        "    # Update metadata\n",
        "    if 'metadata' not in prev_chunk: prev_chunk['metadata'] = {}\n",
        "    prev_chunk['metadata']['chunk_length'] = len(prev_chunk['text'])\n",
        "    prev_chunk['metadata']['page_count'] = len(prev_chunk['page_numbers'])\n",
        "    # Simple merge for other metadata - might need refinement based on specific needs\n",
        "    # For example, how to merge 'systems' or 'parts'? Simple union is often best.\n",
        "    for key in ['systems', 'parts']:\n",
        "        if key in current_chunk.get('metadata', {}):\n",
        "             prev_items = set(prev_chunk['metadata'].get(key, []))\n",
        "             current_items = set(current_chunk['metadata'].get(key, []))\n",
        "             prev_chunk['metadata'][key] = sorted(list(prev_items.union(current_items)))\n",
        "    if current_chunk.get('metadata', {}).get('has_torque'): prev_chunk['metadata']['has_torque'] = True\n",
        "    if current_chunk.get('metadata', {}).get('has_safety'): prev_chunk['metadata']['has_safety'] = True\n",
        "\n",
        "\n",
        "    # Return the last sentences of the *merged* chunk for overlap\n",
        "    merged_sentences = sent_tokenize(prev_chunk['text'])\n",
        "    return merged_sentences[-overlap_sentences:] if merged_sentences else []"
      ],
      "metadata": {
        "id": "k9uiozR_KFT7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "MAIN CHUNKING FUNCTION\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "titlnJlN6kqj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# --- Main Chunking Function ---\n",
        "def create_enriched_chunks(\n",
        "    combined_data: Dict[Tuple[int, float], Dict[str, Any]],\n",
        "    max_chunk_size: int = MAX_CHUNK_SIZE,\n",
        "    min_chunk_size: int = MIN_CHUNK_SIZE,\n",
        "    overlap_sentences: int = OVERLAP_SENTENCE_COUNT,\n",
        "    context_from_previous_batch: Optional[Dict[str, str]] = None,\n",
        "    overlap_sentences_from_previous: Optional[List[str]] = None\n",
        ") -> List[Dict[str, Any]]:\n",
        "    \"\"\"\n",
        "    Creates enriched, contextualized chunks with improved splitting and merging.\n",
        "    - Fixes IndexError.\n",
        "    - Less aggressive splitting before headings.\n",
        "    - More aggressive merging for very small chunks.\n",
        "    \"\"\"\n",
        "    final_chunks = []\n",
        "    if not combined_data:\n",
        "        return final_chunks\n",
        "\n",
        "    sorted_elements = sorted(combined_data.items(), key=lambda item: item[0])\n",
        "\n",
        "    current_chunk_elements = []\n",
        "    current_chunk_text_parts = []\n",
        "    current_size = 0\n",
        "\n",
        "    # Initialize hierarchy from previous batch\n",
        "    current_title = context_from_previous_batch.get('heading_level_1') if context_from_previous_batch else None\n",
        "    current_heading = context_from_previous_batch.get('heading_level_2') if context_from_previous_batch else None\n",
        "    current_subheading = context_from_previous_batch.get('heading_level_3') if context_from_previous_batch else None\n",
        "\n",
        "    # Add overlap from previous batch/chunk if applicable\n",
        "    overlap_text = \"\"\n",
        "    if overlap_sentences_from_previous and overlap_sentences > 0:\n",
        "        overlap_text = \" \".join(overlap_sentences_from_previous)\n",
        "        if overlap_text:\n",
        "            # Use a distinct marker for overlap\n",
        "            current_chunk_text_parts.append(f\"[Overlap Start] {overlap_text} [Overlap End]\\n\\n\")\n",
        "            current_size += len(current_chunk_text_parts[-1])\n",
        "\n",
        "    last_chunk_last_sentences_for_overlap = [] # Track sentences for next overlap\n",
        "\n",
        "    for i, (pos, element) in enumerate(tqdm(sorted_elements, desc=\"Creating chunks\", unit=\"element\")):\n",
        "        element_text = element['text'].strip()\n",
        "        if not element_text: continue # Skip empty elements\n",
        "\n",
        "        element_type = element['type']\n",
        "        element_len = len(element_text)\n",
        "\n",
        "        is_heading = element_type in ['TITLE', 'HEADING', 'SUBHEADING']\n",
        "        is_list_item = element_type == 'LIST_ITEM'\n",
        "        is_table_related = element_type == 'TABLE'\n",
        "\n",
        "        # --- Determine if a new chunk should start BEFORE adding this element ---\n",
        "        new_chunk_needed = False\n",
        "        break_reason = \"\" # For conditional overlap\n",
        "\n",
        "        # 1. Size Limit Check: If adding this element *might* exceed max size\n",
        "        potential_new_size = current_size + element_len + 2 # Estimate with newline/space\n",
        "        if current_size > 0 and potential_new_size > max_chunk_size:\n",
        "            # Only break if current chunk is already substantial\n",
        "            if current_size >= min_chunk_size:\n",
        "                new_chunk_needed = True\n",
        "                break_reason = \"size\"\n",
        "            # If current chunk is small, let it grow (might create oversized chunks if elements are huge)\n",
        "\n",
        "        # 2. Heading Rule (Less Aggressive): Start new chunk *before* a heading,\n",
        "        #    ONLY if current chunk is already large.\n",
        "        if is_heading and current_size >= max_chunk_size * HEADING_SPLIT_THRESHOLD_FACTOR:\n",
        "             new_chunk_needed = True\n",
        "             break_reason = f\"heading_{element_type}\"\n",
        "\n",
        "        # 3. Page Break Rule: Consider breaking, but prioritize semantic flow\n",
        "        if current_chunk_elements and element['page_number'] > current_chunk_elements[-1]['page_number']:\n",
        "            last_element_type = current_chunk_elements[-1]['type']\n",
        "            # Only break on page change if not mid-list/table and chunk is reasonably sized\n",
        "            # Adjust the size threshold here if needed (e.g., min_chunk_size // 2)\n",
        "            if not (is_list_item and last_element_type == 'LIST_ITEM') and \\\n",
        "               not (is_table_related and last_element_type == 'TABLE') and \\\n",
        "               current_size >= min_chunk_size * 0.6: # Break if chunk is >60% of min size on page turn\n",
        "                   new_chunk_needed = True\n",
        "                   break_reason = \"page\"\n",
        "\n",
        "        # --- Finalize the previous chunk if needed ---\\\n",
        "        if new_chunk_needed and current_chunk_elements:\n",
        "            finalized_chunk = _finalize_chunk(\n",
        "                current_chunk_elements,\n",
        "                current_chunk_text_parts,\n",
        "                last_chunk_last_sentences_for_overlap # Pass sentences for potential overlap usage\n",
        "            )\n",
        "\n",
        "            # --- Aggressive Merging Logic ---\n",
        "            merged = False\n",
        "            if final_chunks: # Can only merge if there's a previous chunk\n",
        "                prev_chunk = final_chunks[-1]\n",
        "                finalized_chunk_len = len(finalized_chunk['text'])\n",
        "                prev_chunk_len = len(prev_chunk['text'])\n",
        "                # Approx combined length, actual merge adds separators\n",
        "                combined_len_approx = prev_chunk_len + finalized_chunk_len\n",
        "\n",
        "                # Condition 1: Aggressive merge for VERY small chunks\n",
        "                if finalized_chunk_len < AGGRESSIVE_MERGE_THRESHOLD:\n",
        "                    # Relax context check: only require H1 to match (or both be None)\n",
        "                    h1_compatible = (prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1'))\n",
        "                    # Still respect a reasonable size limit to avoid huge merged chunks\n",
        "                    if h1_compatible and combined_len_approx < max_chunk_size * 1.3: # Allow slightly larger merges\n",
        "                        print(f\"\\nAggressively merging very small chunk ({finalized_chunk_len} chars, Page {finalized_chunk['page_numbers'][0]}) \"\n",
        "                              f\"into previous (Pages {prev_chunk['page_numbers']}) based on H1 context.\")\n",
        "                        # Perform merge and get sentences for next overlap\n",
        "                        last_chunk_last_sentences_for_overlap = _perform_merge(prev_chunk, finalized_chunk, overlap_sentences)\n",
        "                        merged = True\n",
        "                    # else:\n",
        "                        # Optional: Add logging if aggressive merge fails due to size or H1 mismatch\n",
        "                        # print(f\"\\\\nInfo: Very small chunk ({finalized_chunk_len} chars) not merged (H1 mismatch or size limit).\")\n",
        "\n",
        "\n",
        "                # Condition 2: Standard merge (if not aggressively merged and chunk is below min_size)\n",
        "                elif not merged and finalized_chunk_len < min_chunk_size:\n",
        "                     # Existing stricter context check (H1 and H2 must match)\n",
        "                    compatible_context = (\n",
        "                        prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1') and\n",
        "                        prev_chunk.get('heading_level_2') == finalized_chunk.get('heading_level_2')\n",
        "                    )\n",
        "                    if compatible_context and combined_len_approx < max_chunk_size * 1.2: # Standard size limit\n",
        "                         print(f\"\\nMerging small chunk ({finalized_chunk_len} chars, Page {finalized_chunk['page_numbers'][0]}) \"\n",
        "                               f\"into previous (Pages {prev_chunk['page_numbers']}).\")\n",
        "                         # Perform merge and get sentences for next overlap\n",
        "                         last_chunk_last_sentences_for_overlap = _perform_merge(prev_chunk, finalized_chunk, overlap_sentences)\n",
        "                         merged = True\n",
        "                    # else:\n",
        "                         # Optional: Log standard merge failure\n",
        "                         # print(f\"\\\\nWarning: Small chunk ({finalized_chunk_len} chars) not merged (context mismatch or size limit).\")\n",
        "\n",
        "\n",
        "            # Add the finalized chunk to the list IF it wasn't merged\n",
        "            if not merged:\n",
        "                final_chunks.append(finalized_chunk)\n",
        "                # Update last sentences for overlap *only* if a new chunk was added\n",
        "                sentences = sent_tokenize(finalized_chunk['text'])\n",
        "                last_chunk_last_sentences_for_overlap = sentences[-overlap_sentences:] if sentences else []\n",
        "            # If merged, last_chunk_last_sentences_for_overlap was already updated by _perform_merge\n",
        "\n",
        "            # Reset for the new chunk\n",
        "            current_chunk_elements = []\n",
        "            current_chunk_text_parts = []\n",
        "            current_size = 0\n",
        "\n",
        "            # Add overlap conditionally to the *start* of the new chunk\n",
        "            if overlap_sentences > 0 and last_chunk_last_sentences_for_overlap:\n",
        "                 apply_overlap = True\n",
        "                 # Don't add overlap before a major new section starts (TITLE or HEADING)\n",
        "                 if break_reason.startswith(\"heading_\") and element_type in ['TITLE', 'HEADING']:\n",
        "                     apply_overlap = False\n",
        "\n",
        "                 if apply_overlap:\n",
        "                    overlap_text = \" \".join(last_chunk_last_sentences_for_overlap)\n",
        "                    if overlap_text:\n",
        "                        # Use distinct marker\n",
        "                        current_chunk_text_parts.append(f\"[Overlap Start] {overlap_text} [Overlap End]\\n\\n\")\n",
        "                        current_size += len(current_chunk_text_parts[-1])\n",
        "\n",
        "\n",
        "        # Update hierarchy tracking (before adding element to list, but after potential split)\n",
        "        # This reflects the context the *current* element belongs to\n",
        "        if element_type == 'TITLE':\n",
        "            current_title = element_text\n",
        "            current_heading = None\n",
        "            current_subheading = None\n",
        "        elif element_type == 'HEADING':\n",
        "            current_heading = element_text\n",
        "            current_subheading = None\n",
        "        elif element_type == 'SUBHEADING':\n",
        "            current_subheading = element_text\n",
        "        # Store the hierarchy context *with* the element for later use in _finalize_chunk\n",
        "        element['section_hierarchy'] = {1: current_title, 2: current_heading, 3: current_subheading}\n",
        "\n",
        "        # --- Add the current element to the current chunk ---\n",
        "        # Decide prefix BEFORE appending the current element (FIX for IndexError)\n",
        "        prefix = \"\"\n",
        "        if current_chunk_text_parts: # If not the first part of the chunk\n",
        "            # Check the type of the LAST element ADDED to the list\n",
        "            last_element_type = current_chunk_elements[-1]['type'] if current_chunk_elements else None\n",
        "\n",
        "            if is_heading:\n",
        "                 prefix = \"\\n\\n\" # More space before headings\n",
        "            elif is_list_item:\n",
        "                 # Add newline before list item unless the PREVIOUS element was also a list item.\n",
        "                 if last_element_type == 'LIST_ITEM':\n",
        "                      prefix = \"\\n\" # Default to newline for lists for better readability\n",
        "                 else:\n",
        "                      prefix = \"\\n\" # Newline if starting a list or following a non-list item\n",
        "            # Add more sophisticated rules for tables, paragraphs etc. if needed\n",
        "            # elif is_table_related and last_element_type == 'TABLE':\n",
        "            #     prefix = \"\\n\"\n",
        "            else:\n",
        "                 # Add extra newline if previous was a heading\n",
        "                 if last_element_type in ['TITLE', 'HEADING', 'SUBHEADING']:\n",
        "                      prefix = \"\\n\\n\"\n",
        "                 else:\n",
        "                     # Add newline between paragraphs unless previous was list/table?\n",
        "                     if last_element_type not in ['LIST_ITEM', 'TABLE']:\n",
        "                          prefix = \"\\n\\n\" # Default paragraph spacing\n",
        "                     else:\n",
        "                          prefix = \"\\n\"   # Less space after list/table\n",
        "\n",
        "        # Now append the element AFTER determining the prefix based on the previous state\n",
        "        current_chunk_elements.append(element)\n",
        "        current_chunk_text_parts.append(prefix + element_text)\n",
        "        current_size += len(element_text) + len(prefix)\n",
        "\n",
        "    # --- Finalize the very last chunk ---\\\n",
        "    if current_chunk_elements:\n",
        "        finalized_chunk = _finalize_chunk(\n",
        "            current_chunk_elements,\n",
        "            current_chunk_text_parts,\n",
        "            last_chunk_last_sentences_for_overlap\n",
        "        )\n",
        "\n",
        "        # --- Final check for merging the last chunk ---\n",
        "        merged = False\n",
        "        if final_chunks: # Can only merge if there's a previous chunk\n",
        "            prev_chunk = final_chunks[-1]\n",
        "            finalized_chunk_len = len(finalized_chunk['text'])\n",
        "            prev_chunk_len = len(prev_chunk['text'])\n",
        "            combined_len_approx = prev_chunk_len + finalized_chunk_len\n",
        "\n",
        "            # Condition 1: Aggressive merge\n",
        "            if finalized_chunk_len < AGGRESSIVE_MERGE_THRESHOLD:\n",
        "                h1_compatible = (prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1'))\n",
        "                if h1_compatible and combined_len_approx < max_chunk_size * 1.3:\n",
        "                    print(f\"\\nAggressively merging final very small chunk ({finalized_chunk_len} chars, Page {finalized_chunk['page_numbers'][0]}) \"\n",
        "                          f\"into previous (Pages {prev_chunk['page_numbers']}).\")\n",
        "                    _perform_merge(prev_chunk, finalized_chunk, overlap_sentences) # Discard return value here\n",
        "                    merged = True\n",
        "\n",
        "            # Condition 2: Standard merge\n",
        "            elif not merged and finalized_chunk_len < min_chunk_size:\n",
        "                compatible_context = (\n",
        "                    prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1') and\n",
        "                    prev_chunk.get('heading_level_2') == finalized_chunk.get('heading_level_2')\n",
        "                )\n",
        "                if compatible_context and combined_len_approx < max_chunk_size * 1.2:\n",
        "                    print(f\"\\nMerging final small chunk ({finalized_chunk_len} chars, Page {finalized_chunk['page_numbers'][0]}) \"\n",
        "                          f\"into previous (Pages {prev_chunk['page_numbers']}).\")\n",
        "                    _perform_merge(prev_chunk, finalized_chunk, overlap_sentences) # Discard return value here\n",
        "                    merged = True\n",
        "                # else:\n",
        "                    # Optional: Log standard merge failure for the final chunk\n",
        "                    # print(f\"\\\\nWarning: Final chunk is small ({finalized_chunk_len} chars), but not merged.\")\n",
        "\n",
        "        # Add the finalized chunk if it wasn't merged\n",
        "        if not merged:\n",
        "            final_chunks.append(finalized_chunk)\n",
        "\n",
        "    return final_chunks"
      ],
      "metadata": {
        "id": "91ENPNdJ-X-b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# --- Main Chunking Function  ---\n",
        "def create_enriched_chunks(\n",
        "    combined_data: Dict[Tuple[int, float], Dict[str, Any]],\n",
        "    max_chunk_size: int = MAX_CHUNK_SIZE,\n",
        "    min_chunk_size: int = MIN_CHUNK_SIZE,\n",
        "    overlap_sentences: int = OVERLAP_SENTENCE_COUNT,\n",
        "    context_from_previous_batch: Optional[Dict[str, str]] = None,\n",
        "    overlap_sentences_from_previous: Optional[List[str]] = None\n",
        ") -> List[Dict[str, Any]]:\n",
        "    \"\"\"\n",
        "    Creates enriched, contextualized chunks with complete removal of specified terms.\n",
        "    \"\"\"\n",
        "    final_chunks = []\n",
        "    if not combined_data:\n",
        "        return final_chunks\n",
        "\n",
        "    sorted_elements = sorted(combined_data.items(), key=lambda item: item[0])\n",
        "\n",
        "    current_chunk_elements = []\n",
        "    current_chunk_text_parts = []\n",
        "    current_size = 0\n",
        "    current_title = context_from_previous_batch.get('heading_level_1') if context_from_previous_batch else None\n",
        "    current_heading = context_from_previous_batch.get('heading_level_2') if context_from_previous_batch else None\n",
        "    current_subheading = context_from_previous_batch.get('heading_level_3') if context_from_previous_batch else None\n",
        "\n",
        "    overlap_text = \"\"\n",
        "    if overlap_sentences_from_previous and overlap_sentences > 0:\n",
        "        overlap_text = \" \".join(overlap_sentences_from_previous)\n",
        "        if overlap_text:\n",
        "            current_chunk_text_parts.append(f\"[Overlap Start] {overlap_text} [Overlap End]\\n\\n\")\n",
        "            current_size += len(current_chunk_text_parts[-1])\n",
        "\n",
        "    last_chunk_last_sentences_for_overlap = []\n",
        "\n",
        "    for i, (pos, element) in enumerate(tqdm(sorted_elements, desc=\"Creating chunks\", unit=\"element\")):\n",
        "        element_text = element['text'].strip()\n",
        "        if not element_text: continue\n",
        "\n",
        "        element_type = element['type']\n",
        "        element_len = len(element_text)\n",
        "        is_heading = element_type in ['TITLE', 'HEADING', 'SUBHEADING']\n",
        "        is_list_item = element_type == 'LIST_ITEM'\n",
        "        is_table_related = element_type == 'TABLE'\n",
        "\n",
        "        new_chunk_needed = False\n",
        "        break_reason = \"\"\n",
        "        potential_new_size = current_size + element_len + 2\n",
        "        if current_size > 0 and potential_new_size > max_chunk_size:\n",
        "            if current_size >= min_chunk_size:\n",
        "                new_chunk_needed = True\n",
        "                break_reason = \"size\"\n",
        "        if is_heading and current_size >= max_chunk_size * HEADING_SPLIT_THRESHOLD_FACTOR:\n",
        "             new_chunk_needed = True\n",
        "             break_reason = f\"heading_{element_type}\"\n",
        "        if current_chunk_elements and element['page_number'] > current_chunk_elements[-1]['page_number']:\n",
        "            last_element_type = current_chunk_elements[-1]['type']\n",
        "            if not (is_list_item and last_element_type == 'LIST_ITEM') and \\\n",
        "               not (is_table_related and last_element_type == 'TABLE') and \\\n",
        "               current_size >= min_chunk_size * 0.6:\n",
        "                   new_chunk_needed = True\n",
        "                   break_reason = \"page\"\n",
        "\n",
        "        if new_chunk_needed and current_chunk_elements:\n",
        "            finalized_chunk = _finalize_chunk(\n",
        "                current_chunk_elements,\n",
        "                current_chunk_text_parts,\n",
        "                last_chunk_last_sentences_for_overlap\n",
        "            )\n",
        "            merged = False\n",
        "            if final_chunks:\n",
        "                prev_chunk = final_chunks[-1]\n",
        "                finalized_chunk_len = len(finalized_chunk['text'])\n",
        "                prev_chunk_len = len(prev_chunk['text'])\n",
        "                combined_len_approx = prev_chunk_len + finalized_chunk_len\n",
        "                if finalized_chunk_len < AGGRESSIVE_MERGE_THRESHOLD:\n",
        "                    h1_compatible = (prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1'))\n",
        "                    if h1_compatible and combined_len_approx < max_chunk_size * 1.3:\n",
        "                        last_chunk_last_sentences_for_overlap = _perform_merge(prev_chunk, finalized_chunk, overlap_sentences)\n",
        "                        merged = True\n",
        "                elif not merged and finalized_chunk_len < min_chunk_size:\n",
        "                    compatible_context = (\n",
        "                        prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1') and\n",
        "                        prev_chunk.get('heading_level_2') == finalized_chunk.get('heading_level_2')\n",
        "                    )\n",
        "                    if compatible_context and combined_len_approx < max_chunk_size * 1.2:\n",
        "                         last_chunk_last_sentences_for_overlap = _perform_merge(prev_chunk, finalized_chunk, overlap_sentences)\n",
        "                         merged = True\n",
        "            if not merged:\n",
        "                final_chunks.append(finalized_chunk)\n",
        "                sentences = sent_tokenize(finalized_chunk['text'])\n",
        "                last_chunk_last_sentences_for_overlap = sentences[-overlap_sentences:] if sentences else []\n",
        "            current_chunk_elements = []\n",
        "            current_chunk_text_parts = []\n",
        "            current_size = 0\n",
        "            if overlap_sentences > 0 and last_chunk_last_sentences_for_overlap:\n",
        "                 apply_overlap = True\n",
        "                 if break_reason.startswith(\"heading_\") and element_type in ['TITLE', 'HEADING']:\n",
        "                     apply_overlap = False\n",
        "                 if apply_overlap:\n",
        "                    overlap_text = \" \".join(last_chunk_last_sentences_for_overlap)\n",
        "                    if overlap_text:\n",
        "                        current_chunk_text_parts.append(f\"[Overlap Start] {overlap_text} [Overlap End]\\n\\n\")\n",
        "                        current_size += len(current_chunk_text_parts[-1])\n",
        "\n",
        "        if element_type == 'TITLE':\n",
        "            current_title = element_text\n",
        "            current_heading = None\n",
        "            current_subheading = None\n",
        "        elif element_type == 'HEADING':\n",
        "            current_heading = element_text\n",
        "            current_subheading = None\n",
        "        elif element_type == 'SUBHEADING':\n",
        "            current_subheading = element_text\n",
        "        element['section_hierarchy'] = {1: current_title, 2: current_heading, 3: current_subheading}\n",
        "\n",
        "\n",
        "        prefix = \"\"\n",
        "        if current_chunk_text_parts:\n",
        "            last_element_type = current_chunk_elements[-1]['type'] if current_chunk_elements else None\n",
        "            if is_heading:\n",
        "                 prefix = \"\\n\\n\"\n",
        "            elif is_list_item:\n",
        "                 if last_element_type == 'LIST_ITEM':\n",
        "                      prefix = \"\\n\"\n",
        "                 else:\n",
        "                      prefix = \"\\n\"\n",
        "            else:\n",
        "                 if last_element_type in ['TITLE', 'HEADING', 'SUBHEADING']:\n",
        "                      prefix = \"\\n\\n\"\n",
        "                 else:\n",
        "                     if last_element_type not in ['LIST_ITEM', 'TABLE']:\n",
        "                          prefix = \"\\n\\n\"\n",
        "                     else:\n",
        "                          prefix = \"\\n\"\n",
        "\n",
        "        current_chunk_elements.append(element)\n",
        "        current_chunk_text_parts.append(prefix + element_text)\n",
        "        current_size += len(element_text) + len(prefix)\n",
        "\n",
        "    if current_chunk_elements:\n",
        "        finalized_chunk = _finalize_chunk(\n",
        "            current_chunk_elements,\n",
        "            current_chunk_text_parts,\n",
        "            last_chunk_last_sentences_for_overlap\n",
        "        )\n",
        "        merged = False\n",
        "        if final_chunks:\n",
        "            prev_chunk = final_chunks[-1]\n",
        "            finalized_chunk_len = len(finalized_chunk['text'])\n",
        "            prev_chunk_len = len(prev_chunk['text'])\n",
        "            combined_len_approx = prev_chunk_len + finalized_chunk_len\n",
        "            if finalized_chunk_len < AGGRESSIVE_MERGE_THRESHOLD:\n",
        "                h1_compatible = (prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1'))\n",
        "                if h1_compatible and combined_len_approx < max_chunk_size * 1.3:\n",
        "                    _perform_merge(prev_chunk, finalized_chunk, overlap_sentences)\n",
        "                    merged = True\n",
        "            elif not merged and finalized_chunk_len < min_chunk_size:\n",
        "                compatible_context = (\n",
        "                    prev_chunk.get('heading_level_1') == finalized_chunk.get('heading_level_1') and\n",
        "                    prev_chunk.get('heading_level_2') == finalized_chunk.get('heading_level_2')\n",
        "                )\n",
        "                if compatible_context and combined_len_approx < max_chunk_size * 1.2:\n",
        "                    _perform_merge(prev_chunk, finalized_chunk, overlap_sentences)\n",
        "                    merged = True\n",
        "        if not merged:\n",
        "            final_chunks.append(finalized_chunk)\n",
        "\n",
        "    return final_chunks"
      ],
      "metadata": {
        "id": "WkhXbdXpbfMy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "ENAHNCED FEATURE EXTRACTION\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "PcUmD7ZH-sdg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def _finalize_chunk(\n",
        "    elements: List[Dict[str, Any]],\n",
        "    text_parts: List[str],\n",
        "    last_chunk_sentences_for_overlap: List[str] # Used only if overlap logic needs it here\n",
        ") -> Dict[str, Any]:\n",
        "    \"\"\"Helper to finalize a chunk: combine text, extract features, format, handle warnings.\"\"\"\n",
        "    # Join text parts respecting prefixes added earlier\n",
        "    raw_text = \"\".join(text_parts).strip()\n",
        "\n",
        "    # --- Centralized Safety Warning Handling ---\n",
        "    warnings_found = []\n",
        "    processed_text = raw_text\n",
        "    # Use finditer to get match objects for accurate replacement\n",
        "    for match in WARNING_PATTERN.finditer(raw_text):\n",
        "        warning_type = match.group(1).upper()\n",
        "        # Group 2 is the main content, Group 3 is the terminator (.!? or \\n\\n)\n",
        "        content = (match.group(2) + match.group(3)).strip()\n",
        "        full_match_text = match.group(0) # The entire matched string\n",
        "\n",
        "        # Check if we've already processed this exact text span to avoid duplicates from overlapping matches\n",
        "        already_added = False\n",
        "        for _, existing_content, _ in warnings_found:\n",
        "            if content in existing_content or existing_content in content:\n",
        "                 # Basic check to avoid adding subsections of already captured warnings\n",
        "                 # This might need refinement depending on pattern behavior\n",
        "                 already_added = True\n",
        "                 break\n",
        "        if not already_added:\n",
        "             warnings_found.append((warning_type, content, full_match_text))\n",
        "\n",
        "\n",
        "    warning_block = \"\"\n",
        "    if warnings_found:\n",
        "        warning_texts = [f\"⚠️ {wtype}: {wcontent}\" for wtype, wcontent, _ in warnings_found]\n",
        "        warning_block = \"\\n\".join(warning_texts) + \"\\n\\n\"\n",
        "\n",
        "        # Remove original warnings from the text - be careful!\n",
        "        # Sort by length descending to remove longer matches first (helps with nested cases)\n",
        "        warnings_found.sort(key=lambda x: len(x[2]), reverse=True)\n",
        "        temp_text = processed_text\n",
        "        for _, _, full_match_text in warnings_found:\n",
        "             # Replace only the first occurrence found in this pass to avoid cascading issues\n",
        "             temp_text = temp_text.replace(full_match_text, \"\", 1)\n",
        "        processed_text = temp_text\n",
        "\n",
        "\n",
        "    # Combine warning block and cleaned text\n",
        "    final_text = warning_block + processed_text\n",
        "    final_text = re.sub(r'\\n\\s*\\n', '\\n\\n', final_text) # Clean up multiple blank lines\n",
        "    final_text = re.sub(r'[ \\t]+', ' ', final_text) # Consolidate spaces/tabs\n",
        "    final_text = final_text.strip()\n",
        "\n",
        "    # Determine context from hierarchy (get the most specific available)\n",
        "    final_hierarchy = {}\n",
        "    current_title, current_heading, current_subheading = None, None, None\n",
        "    for element in elements:\n",
        "        h = element.get('section_hierarchy', {})\n",
        "        # Update if the element provides a non-empty value for a level\n",
        "        if h.get(1): current_title = h[1]\n",
        "        if h.get(2): current_heading = h[2]\n",
        "        if h.get(3): current_subheading = h[3]\n",
        "    # Store the final determined hierarchy for the chunk\n",
        "    final_hierarchy = {1: current_title, 2: current_heading, 3: current_subheading}\n",
        "\n",
        "    # Create breadcrumb from hierarchy\n",
        "    context_parts = [final_hierarchy.get(i) for i in sorted(final_hierarchy.keys()) if final_hierarchy.get(i)]\n",
        "    breadcrumb_trail = ' > '.join(context_parts)\n",
        "\n",
        "    # Extract features from the combined text\n",
        "    features = extract_chunk_features(final_text) # Run on final text including warnings\n",
        "\n",
        "\n",
        "    page_numbers = sorted(list(set(elem['page_number'] for elem in elements)))\n",
        "\n",
        "    chunk = {\n",
        "        'text': final_text,\n",
        "        'context': breadcrumb_trail, # For simpler access\n",
        "        'page_numbers': page_numbers,\n",
        "        'breadcrumb_trail': breadcrumb_trail, # Explicitly keep\n",
        "        'heading_level_1': final_hierarchy.get(1),\n",
        "        'heading_level_2': final_hierarchy.get(2),\n",
        "        'heading_level_3': final_hierarchy.get(3),\n",
        "        'has_overlap_prefix': text_parts[0].startswith(\"[Overlap]\") if text_parts else False\n",
        "    }\n",
        "\n",
        "    chunk.update(features) # Add all extracted features\n",
        "\n",
        "    # Add content type indicators based on features\n",
        "    content_type = set()\n",
        "    if chunk.get('is_procedure', False): content_type.add('procedure')\n",
        "    if chunk.get('contains_table', False): content_type.add('table')\n",
        "    if chunk.get('figure_references'): content_type.add('figure')\n",
        "    if chunk.get('safety_notices'): content_type.add('safety') # Populated by extract_chunk_features\n",
        "    if chunk.get('is_diagnostic_content', False): content_type.add('diagnostic')\n",
        "    if chunk.get('torque_specifications'): content_type.add('specification')\n",
        "    if chunk.get('maintenance_intervals'): content_type.add('maintenance')\n",
        "    if chunk.get('contains_diagnostic_codes'): content_type.add('dtc')\n",
        "    if content_type:\n",
        "        chunk['content_type'] = sorted(list(content_type))\n",
        "\n",
        "    # Consolidate Metadata\n",
        "    metadata = {\n",
        "        'source_pages': f\"{min(page_numbers)}-{max(page_numbers)}\" if page_numbers else \"N/A\",\n",
        "        'page_count': len(page_numbers),\n",
        "        'chunk_length': len(final_text),\n",
        "    }\n",
        "    if chunk.get('vehicle_systems'): metadata['systems'] = chunk['vehicle_systems']\n",
        "    if chunk.get('torque_specifications'): metadata['has_torque'] = True\n",
        "    if chunk.get('safety_notices'): metadata['has_safety'] = True\n",
        "\n",
        "    chunk['metadata'] = metadata\n",
        "\n",
        "    # Clean up redundant top-level keys if they exist in metadata\n",
        "    for key in ['vehicle_systems', 'torque_specifications', 'safety_notices']:\n",
        "       if key in chunk:\n",
        "            del chunk[key]\n",
        "\n",
        "    return chunk\n",
        "\n",
        "# --- Enhanced Feature Extraction Function ---\n",
        "# (Largely similar, ensures patterns match against potentially modified final text)\n",
        "def extract_chunk_features(text: str) -> Dict[str, Any]:\n",
        "    \"\"\"Extracts important features from chunk text with Ford-specific enhancements.\"\"\"\n",
        "    features = {}\n",
        "\n",
        "    # Safety notices (relies on the prioritized format from _finalize_chunk)\n",
        "    # Extract from the \"⚠️ TYPE: content\" format\n",
        "    safety_matches = re.findall(r\"⚠️\\s*(WARNING|CAUTION|DANGER|IMPORTANT|NOTE|ATTENTION):\\s*(.*?)(?=\\n⚠️|\\Z)\", text, re.DOTALL)\n",
        "    if safety_matches:\n",
        "        features['safety_notices'] = [\n",
        "            {'type': match[0].upper(), 'content': match[1].strip()}\n",
        "            for match in safety_matches\n",
        "        ]\n",
        "\n",
        "    # Procedural steps\n",
        "    # Simplified: Check for common list markers or step indicators\n",
        "    step_indicators = re.findall(r'(?m)^\\s*(\\d+\\.|\\*|\\-|•)\\s+', text) # Check for lines starting with markers\n",
        "    numbered_steps = re.search(r'\\b(step\\s+\\d+)\\b', text, re.IGNORECASE)\n",
        "    if step_indicators or numbered_steps:\n",
        "         features['is_procedure'] = True # Flag existence, detailed steps maybe too complex here\n",
        "         # Example: Extract first few steps if needed\n",
        "         # procedural_steps_pattern = re.compile(r'(?m)^\\s*(\\d+\\.|\\*|\\-|•)\\s+(.*)')\n",
        "         # steps = procedural_steps_pattern.findall(text)\n",
        "         # if steps: features['procedural_steps_preview'] = [f\"{s[0]} {s[1][:100]}...\" for s in steps[:3]]\n",
        "\n",
        "\n",
        "    # Torque specifications\n",
        "    torque_pattern = re.compile(r'(\\d+(?:\\.\\d+)?)\\s*(?:Nm|N·m|lb-ft|lb·ft|ft-lb|ft·lb|lb-in|lb·in|in-lb|in·lb)', re.IGNORECASE)\n",
        "    torque_specs = torque_pattern.findall(text)\n",
        "    if torque_specs:\n",
        "        # Simple extraction for now, units require more parsing\n",
        "        features['torque_specifications'] = [{'value': spec} for spec in torque_specs]\n",
        "\n",
        "\n",
        "    # Table detection (basic heuristic)\n",
        "    # More than 2 lines with multiple pipe symbols or consistent spacing might indicate a table\n",
        "    lines = text.splitlines()\n",
        "    pipe_lines = [line for line in lines if '|' in line.strip()]\n",
        "    plus_lines = [line for line in lines if '+' in line.strip() and '-' in line.strip()]\n",
        "    if len(pipe_lines) > 2 or len(plus_lines) > 1 :\n",
        "        features['contains_table'] = True\n",
        "\n",
        "    # Figure references\n",
        "    figure_pattern = re.compile(r'\\b(Fig(?:ure)?\\.?\\s+\\d+[a-zA-Z]?)\\b', re.IGNORECASE)\n",
        "    figures = figure_pattern.findall(text)\n",
        "    if figures:\n",
        "        features['figure_references'] = sorted(list(set(figures)))\n",
        "\n",
        "    # Vehicle systems (simple keyword check)\n",
        "    systems_dict = { # Keep this relatively concise\n",
        "        'engine': ['engine', 'ecoboost', 'powerstroke', 'cylinder', 'piston', 'timing'],\n",
        "        'transmission': ['transmission', 'gearbox', 'clutch', 'transaxle', 'powershift', 'atf'],\n",
        "        'electrical': ['battery', 'alternator', 'starter', 'fuse', 'relay', 'wiring', 'module', 'sensor', 'pcm', 'bcm'],\n",
        "        'fuel': ['fuel tank', 'fuel pump', 'fuel filter', 'injector', 'emission'],\n",
        "        'cooling': ['coolant', 'radiator', 'thermostat', 'water pump'],\n",
        "        'brakes': ['brake', 'abs', 'rotor', 'caliper', 'pad'],\n",
        "        'suspension': ['suspension', 'shock', 'strut', 'spring', 'control arm', 'alignment'],\n",
        "        'drivetrain': ['transfer case', '4x4', 'awd', 'differential', 'axle', 'driveshaft'],\n",
        "        'steering': ['steering', 'epas', 'rack', 'pinion', 'tie rod'],\n",
        "        'body': ['door', 'hood', 'trunk', 'liftgate', 'bumper', 'windshield', 'wiper', 'mirror', 'body control module'],\n",
        "        'hvac': ['climate', 'a/c', 'heater', 'compressor', 'condenser'],\n",
        "        'exhaust':['exhaust', 'muffler', 'catalytic converter']\n",
        "    }\n",
        "    detected_systems = set()\n",
        "    text_lower = text.lower()\n",
        "    for system, keywords in systems_dict.items():\n",
        "        for keyword in keywords:\n",
        "            if re.search(r'\\b' + keyword + r'\\b', text_lower):\n",
        "                detected_systems.add(system)\n",
        "                break\n",
        "    if detected_systems:\n",
        "        features['vehicle_systems'] = sorted(list(detected_systems))\n",
        "\n",
        "    # Maintenance intervals\n",
        "    interval_pattern = re.compile(r'(\\d{1,3}(?:,\\d{3})*|\\d+)\\s+(miles|mi|kilometers|km|months|years)', re.IGNORECASE)\n",
        "    intervals = interval_pattern.findall(text)\n",
        "    if intervals:\n",
        "        features['maintenance_intervals'] = [\n",
        "            {'value': match[0].replace(',', ''), 'unit': match[1].lower()}\n",
        "            for match in intervals\n",
        "        ]\n",
        "\n",
        "    # Diagnostic codes (OBDII)\n",
        "    obd_code_pattern = re.compile(r'\\b([PBUC]\\d{4})\\b')\n",
        "    dtcs = obd_code_pattern.findall(text)\n",
        "    if dtcs:\n",
        "        features['contains_diagnostic_codes'] = True\n",
        "        features['diagnostic_codes'] = sorted(list(set(dtcs)))\n",
        "\n",
        "    diagnostic_keywords = re.search(r'\\b(diagnos(e|is|tic)|troubleshoot|symptom|fault|error|dtc|code|check engine)\\b', text, re.IGNORECASE)\n",
        "    if diagnostic_keywords:\n",
        "        features['is_diagnostic_content'] = True\n",
        "\n",
        "    return features"
      ],
      "metadata": {
        "id": "5EmEeRuh-jfD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "POST PROCESSING\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "mD75dldR-3AE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# --- Post Processing Function ---\n",
        "def post_process_chunks(chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:\n",
        "    \"\"\"Perform final post-processing: IDs, links, related chunks.\"\"\"\n",
        "    if not chunks: return []\n",
        "\n",
        "    # Sort chunks by first page number for stable ID assignment\n",
        "    # Removed 'id' from sort key as it doesn't exist yet\n",
        "    chunks.sort(key=lambda x: (min(x['page_numbers']) if x['page_numbers'] else 0))\n",
        "\n",
        "    # --- Pass 1: Assign IDs and Link Previous ---\n",
        "    for i, chunk in enumerate(chunks):\n",
        "        # Assign ID to the current chunk\n",
        "        chunk['id'] = f\"chunk_{i+1:05d}\" # Pad more for potentially many chunks\n",
        "\n",
        "        # Link current chunk to the previous chunk (if it exists)\n",
        "        if i > 0:\n",
        "            chunk['prev_chunk_id'] = chunks[i-1]['id']\n",
        "\n",
        "    # --- Pass 2: Link Next and Calculate Related Chunks ---\n",
        "    heading_map_h2 = {}\n",
        "    heading_map_h3 = {}\n",
        "    feature_map = {}\n",
        "\n",
        "    # Build maps using the now assigned IDs and link next chunks\n",
        "    for i, chunk in enumerate(chunks):\n",
        "        chunk_id = chunk['id'] # ID now exists\n",
        "\n",
        "        # Link current chunk to the next chunk (if it exists)\n",
        "        if i < len(chunks) - 1:\n",
        "             chunk['next_chunk_id'] = chunks[i+1]['id']\n",
        "             # Ensure the previous chunk also has its next_chunk_id set correctly\n",
        "             # (This covers the case where we link from the previous loop)\n",
        "             # If chunk i-1 exists, its next_chunk_id should point to chunk i\n",
        "             if i > 0 and 'next_chunk_id' not in chunks[i-1]:\n",
        "                 chunks[i-1]['next_chunk_id'] = chunk_id\n",
        "\n",
        "\n",
        "        # Build maps for related chunk calculation\n",
        "        h2 = chunk.get('heading_level_2')\n",
        "        h3 = chunk.get('heading_level_3')\n",
        "        if h3: # Prioritize H3 for specificity\n",
        "             heading_map_h3.setdefault(h3, []).append(chunk_id)\n",
        "        elif h2:\n",
        "             heading_map_h2.setdefault(h2, []).append(chunk_id)\n",
        "\n",
        "\n",
        "    # Assign related chunks (using the maps built above)\n",
        "    for chunk in chunks:\n",
        "        related_ids = set()\n",
        "        chunk_id = chunk['id']\n",
        "\n",
        "        # Related by heading\n",
        "        h3 = chunk.get('heading_level_3')\n",
        "        h2 = chunk.get('heading_level_2')\n",
        "        if h3 and h3 in heading_map_h3:\n",
        "             related_ids.update(cid for cid in heading_map_h3[h3] if cid != chunk_id)\n",
        "        elif h2 and h2 in heading_map_h2: # Use H2 if H3 isn't specific enough or absent\n",
        "             related_ids.update(cid for cid in heading_map_h2[h2] if cid != chunk_id)\n",
        "\n",
        "        if related_ids:\n",
        "             # Convert to list, sort, limit\n",
        "             chunk['related_chunks'] = sorted(list(related_ids))[:5] # Limit for clarity\n",
        "\n",
        "    return chunks"
      ],
      "metadata": {
        "id": "OHzKyjki-qhV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "VALIDATION\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "N8Ra9MaA-79a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# --- Validation Function ---\n",
        "def validate_chunking(\n",
        "    final_chunks: List[Dict[str, Any]],\n",
        "    min_size: int = MIN_CHUNK_SIZE,\n",
        "    max_size: int = MAX_CHUNK_SIZE,\n",
        "    overlap_expected: bool = (OVERLAP_SENTENCE_COUNT > 0)\n",
        "    ):\n",
        "    \"\"\"Validates the generated chunks based on size, overlap, and context.\"\"\"\n",
        "    print(\"\\n--- Starting Chunk Validation ---\")\n",
        "    num_chunks = len(final_chunks)\n",
        "    if num_chunks == 0:\n",
        "        print(\"Validation Error: No chunks were generated.\")\n",
        "        return\n",
        "\n",
        "    print(f\"Total Chunks Generated: {num_chunks}\")\n",
        "\n",
        "    sizes = [len(chunk['text']) for chunk in final_chunks]\n",
        "    below_min = sum(1 for s in sizes if s < min_size)\n",
        "    # Allow some flexibility above max size due to sentence splitting logic\n",
        "    oversized_threshold = max_size * 1.20\n",
        "    above_max = sum(1 for s in sizes if s > max_size)\n",
        "    significantly_above_max = sum(1 for s in sizes if s > oversized_threshold)\n",
        "    avg_size = sum(sizes) / num_chunks if num_chunks > 0 else 0\n",
        "\n",
        "    print(\"\\nChunk Size Analysis:\")\n",
        "    print(f\"  Average Size: {avg_size:.0f} characters\")\n",
        "    print(f\"  Min Size: {min(sizes)} | Max Size: {max(sizes)}\")\n",
        "    print(f\"  Chunks below Min ({min_size}): {below_min} ({below_min/num_chunks:.1%})\")\n",
        "    print(f\"  Chunks above Max ({max_size}): {above_max} ({above_max/num_chunks:.1%})\")\n",
        "    print(f\"  Chunks significantly above Max ({oversized_threshold:.0f}): {significantly_above_max} ({significantly_above_max/num_chunks:.1%})\")\n",
        "\n",
        "    if below_min > num_chunks * 0.02: # Tolerate fewer small chunks due to merging\n",
        "         print(f\"  WARNING: Still found {below_min} chunks below {min_size} chars. Review merge logic or MIN_CHUNK_SIZE.\")\n",
        "    if significantly_above_max > num_chunks * 0.05: # If > 5% are significantly oversized\n",
        "         print(f\"  WARNING: {significantly_above_max} chunks significantly exceed max size ({oversized_threshold:.0f}). Review splitting logic or MAX_CHUNK_SIZE.\")\n",
        "\n",
        "    # Overlap Validation\n",
        "    if overlap_expected:\n",
        "        overlap_chunks = [chunk for chunk in final_chunks if chunk.get('has_overlap_prefix')]\n",
        "        num_overlap = len(overlap_chunks)\n",
        "        expected_overlap = num_chunks - 1 if num_chunks > 0 else 0 # Rough expectation\n",
        "        print(f\"\\nOverlap Analysis (Expected sentences: {OVERLAP_SENTENCE_COUNT}):\")\n",
        "        print(f\"  Chunks with Overlap Prefix: {num_overlap} ({num_overlap/num_chunks:.1%})\")\n",
        "        # This check is less reliable with conditional overlap\n",
        "        # if num_overlap < expected_overlap * 0.7 and num_chunks > 1:\n",
        "        #      print(\"  INFO: Fewer chunks have overlap than expected, possibly due to conditional logic (breaks at headings).\")\n",
        "    else:\n",
        "         print(\"\\nOverlap Analysis: Overlap Disabled.\")\n",
        "\n",
        "    # Context Coherence Check (Sample) - check hierarchy levels\n",
        "    print(\"\\nContext Coherence Check (Sample):\")\n",
        "    coherent = True\n",
        "    for i in range(1, min(num_chunks, 10)): # Check first few transitions\n",
        "        prev_chunk = final_chunks[i-1]\n",
        "        curr_chunk = final_chunks[i]\n",
        "\n",
        "        # Check if context seems logically consistent (same level or one level down/up)\n",
        "        p_h1, p_h2, p_h3 = prev_chunk.get('heading_level_1'), prev_chunk.get('heading_level_2'), prev_chunk.get('heading_level_3')\n",
        "        c_h1, c_h2, c_h3 = curr_chunk.get('heading_level_1'), curr_chunk.get('heading_level_2'), curr_chunk.get('heading_level_3')\n",
        "\n",
        "        context_jump = False\n",
        "        if p_h1 != c_h1: context_jump = True # Major jump if H1 changes\n",
        "        elif p_h2 != c_h2:\n",
        "             # Allow change if H3 is None in both (moving between H2 sections)\n",
        "             if p_h3 is not None or c_h3 is not None:\n",
        "                  # Allow moving up from H3 to parent H2\n",
        "                  if not (p_h3 is not None and c_h3 is None and p_h2 == c_h2):\n",
        "                      context_jump = True\n",
        "        # Allow H3 to change if H1 and H2 are the same\n",
        "        # This basic check is not perfect but flags major discontinuities\n",
        "\n",
        "        if context_jump:\n",
        "             print(f\"  Potential Context Jump:\")\n",
        "             print(f\"    Chunk {i} Context: {prev_chunk.get('breadcrumb_trail', 'N/A')}\")\n",
        "             print(f\"    Chunk {i+1} Context: {curr_chunk.get('breadcrumb_trail', 'N/A')}\")\n",
        "             coherent = False\n",
        "             # break\n",
        "\n",
        "    if coherent:\n",
        "         print(\"  Context appears generally coherent in sample based on hierarchy.\")\n",
        "\n",
        "    # Readability Check Prompt\n",
        "    print(\"\\nReadability Check:\")\n",
        "    print(\"  Please manually review chunks in the generated HTML file ('enriched_chunks_viewer.html'),\")\n",
        "    print(\"  especially those containing lists, procedures, tables, or warnings,\")\n",
        "    print(\"  to ensure text flows correctly and formatting is acceptable.\")\n",
        "\n",
        "    print(\"\\n--- Validation Complete ---\")"
      ],
      "metadata": {
        "id": "Cz0jGFPi-1MA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "CHUNK HTML EXPORT\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "VAgdVfuA_IvM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import html # Import the html module for escaping\n",
        "\n",
        "# --- HTML Export Function ---\n",
        "def create_html_viewer(chunks: List[Dict[str, Any]], filename: str = \"enriched_chunks_viewer.html\"):\n",
        "    \"\"\"Generates an HTML file to browse the chunks.\"\"\"\n",
        "    if not chunks:\n",
        "        print(\"No chunks to generate HTML viewer for.\")\n",
        "        return None\n",
        "\n",
        "    num_chunks = len(chunks)\n",
        "\n",
        "    # Use f-string for the main structure and escape CSS braces {{ }}\n",
        "    html_content = f\"\"\"\n",
        "<!DOCTYPE html>\n",
        "<html lang=\"en\">\n",
        "<head>\n",
        "    <meta charset=\"UTF-8\">\n",
        "    <meta name=\"viewport\" content=\"width=device-width, initial-scale=1.0\">\n",
        "    <title>Chunk Viewer</title>\n",
        "    <style>\n",
        "        /* Main styling */\n",
        "        body {{ font-family: sans-serif; line-height: 1.5; padding: 20px; }}\n",
        "        .widget-label {{\n",
        "            min-width: 160px;\n",
        "            max-width: 160px;\n",
        "            font-weight: bold;\n",
        "            color: #444;\n",
        "        }}\n",
        "\n",
        "        /* Header styling */\n",
        "        .search-header {{\n",
        "            background: linear-gradient(to right, #3a7bd5, #00d2ff);\n",
        "            color: white;\n",
        "            padding: 15px;\n",
        "            border-radius: 8px;\n",
        "            margin-bottom: 20px;\n",
        "            box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);\n",
        "            text-align: center;\n",
        "        }}\n",
        "\n",
        "        .search-header h2 {{\n",
        "            margin-top: 0;\n",
        "            margin-bottom: 10px;\n",
        "        }}\n",
        "\n",
        "        /* Panel styling */\n",
        "        .panel {{\n",
        "            border-radius: 8px;\n",
        "            padding: 15px;\n",
        "            margin-bottom: 15px;\n",
        "            box-shadow: 0 2px 5px rgba(0, 0, 0, 0.05);\n",
        "            border: 1px solid #e0e0e0;\n",
        "            background-color: #fff;\n",
        "        }}\n",
        "\n",
        "        .panel-header {{\n",
        "            font-size: 16px;\n",
        "            font-weight: bold;\n",
        "            margin-bottom: 10px;\n",
        "            color: #2c3e50;\n",
        "            border-bottom: 2px solid #ecf0f1;\n",
        "            padding-bottom: 8px;\n",
        "        }}\n",
        "\n",
        "        /* Results styling */\n",
        "        .result-item {{\n",
        "            background-color: #f8f9fa;\n",
        "            border-left: 4px solid #3a7bd5;\n",
        "            border-radius: 5px;\n",
        "            padding: 15px;\n",
        "            margin-bottom: 15px;\n",
        "            transition: all 0.2s ease;\n",
        "            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);\n",
        "            page-break-inside: avoid; /* Try to prevent printing breaks inside */\n",
        "        }}\n",
        "\n",
        "        .result-item:hover {{\n",
        "            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);\n",
        "            transform: translateY(-2px);\n",
        "        }}\n",
        "\n",
        "        .result-item h3 {{\n",
        "            color: #2c3e50;\n",
        "            margin-top: 0;\n",
        "            font-size: 18px;\n",
        "            border-bottom: 1px solid #eee;\n",
        "            padding-bottom: 8px;\n",
        "        }}\n",
        "        .chunk-header {{\n",
        "            border-bottom: 1px solid #eee;\n",
        "            margin-bottom: 10px;\n",
        "            padding-bottom: 10px;\n",
        "            display: flex;\n",
        "            justify-content: space-between;\n",
        "            align-items: center;\n",
        "        }}\n",
        "        .chunk-id {{ font-weight: bold; font-size: 1.1em; color: #333; }}\n",
        "        .chunk-nav a {{ text-decoration: none; margin-left: 10px; color: #007bff; }}\n",
        "        .chunk-nav a:hover {{ text-decoration: underline; }}\n",
        "        .chunk-meta {{ font-size: 0.9em; color: #555; margin-bottom: 10px; }}\n",
        "        .chunk-meta strong {{ color: #333; }}\n",
        "        .chunk-text {{ white-space: pre-wrap; word-wrap: break-word; }}\n",
        "        .features {{ margin-top: 10px; font-size: 0.9em; }}\n",
        "        .features strong {{ display: block; margin-bottom: 5px; color: #444; }}\n",
        "        .features ul {{ margin: 0; padding-left: 20px; list-style: disc; }}\n",
        "        .features li {{ margin-bottom: 3px; }}\n",
        "        .warning {{ color: #d9534f; font-weight: bold; }} /* Style for ⚠️ */\n",
        "        .overlap-prefix {{ color: #777; font-style: italic; font-size: 0.9em; }}\n",
        "\n",
        "        /* Score visualization */\n",
        "        .score-bar {{\n",
        "            height: 8px;\n",
        "            background-color: #ecf0f1;\n",
        "            border-radius: 4px;\n",
        "            margin: 5px 0;\n",
        "            overflow: hidden;\n",
        "        }}\n",
        "\n",
        "        .score-fill {{\n",
        "            height: 100%;\n",
        "            background: linear-gradient(to right, #3a7bd5, #00d2ff);\n",
        "        }}\n",
        "\n",
        "        .text-score-fill {{\n",
        "            height: 100%;\n",
        "            background: linear-gradient(to right, #11998e, #38ef7d);\n",
        "        }}\n",
        "\n",
        "        /* Badges for metadata */\n",
        "        .badge {{\n",
        "            display: inline-block;\n",
        "            padding: 3px 8px;\n",
        "            border-radius: 12px;\n",
        "            font-size: 12px;\n",
        "            font-weight: bold;\n",
        "            margin-right: 5px;\n",
        "            color: white;\n",
        "        }}\n",
        "\n",
        "        .page-badge {{\n",
        "            background-color: #3498db;\n",
        "        }}\n",
        "\n",
        "        .type-badge {{\n",
        "            background-color: #9b59b6;\n",
        "        }}\n",
        "\n",
        "        /* Content panels */\n",
        "        .context-panel {{\n",
        "            background-color: #f1f8fb;\n",
        "            border-radius: 5px;\n",
        "            padding: 8px 12px;\n",
        "            margin: 8px 0;\n",
        "            border-left: 3px solid #3498db;\n",
        "        }}\n",
        "\n",
        "        .text-panel {{\n",
        "            max-height: 200px;\n",
        "            overflow-y: auto;\n",
        "            background-color: #f9f9f9;\n",
        "            border-radius: 5px;\n",
        "            padding: 12px;\n",
        "            border: 1px solid #eaeaea;\n",
        "            font-family: system-ui, -apple-system, sans-serif;\n",
        "            line-height: 1.6;\n",
        "        }}\n",
        "\n",
        "        /* Button styling */\n",
        "        .search-button {{\n",
        "            text-transform: uppercase;\n",
        "            letter-spacing: 1px;\n",
        "            transition: all 0.3s ease;\n",
        "        }}\n",
        "\n",
        "        .search-button:hover {{\n",
        "            transform: translateY(-2px);\n",
        "            box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);\n",
        "        }}\n",
        "\n",
        "        /* Loading animation */\n",
        "        .spinner {{\n",
        "            border: 4px solid rgba(0, 0, 0, 0.1);\n",
        "            width: 36px;\n",
        "            height: 36px;\n",
        "            border-radius: 50%;\n",
        "            border-left-color: #3a7bd5;\n",
        "            animation: spin 1s linear infinite;\n",
        "            margin: 20px auto;\n",
        "        }}\n",
        "\n",
        "        @keyframes spin {{\n",
        "            0% {{ transform: rotate(0deg); }}\n",
        "            100% {{ transform: rotate(360deg); }}\n",
        "        }}\n",
        "\n",
        "        /* Divider styling */\n",
        "        .divider {{\n",
        "            display: flex;\n",
        "            align-items: center;\n",
        "            margin: 15px 0;\n",
        "        }}\n",
        "\n",
        "        .divider span {{\n",
        "            padding: 0 10px;\n",
        "            color: #718093;\n",
        "            font-weight: bold;\n",
        "        }}\n",
        "\n",
        "        .divider:before,\n",
        "        .divider:after {{\n",
        "            content: \"\";\n",
        "            flex: 1;\n",
        "            border-bottom: 1px solid #ecf0f1;\n",
        "        }}\n",
        "\n",
        "        /* Status area */\n",
        "        .status-box {{\n",
        "            background-color: #f8f9fa;\n",
        "            border-radius: 5px;\n",
        "            padding: 12px;\n",
        "            margin: 10px 0;\n",
        "            border-left: 4px solid #3498db;\n",
        "        }}\n",
        "\n",
        "        /* No results */\n",
        "        .no-results {{\n",
        "            text-align: center;\n",
        "            padding: 40px;\n",
        "            color: #666;\n",
        "        }}\n",
        "\n",
        "        .no-results-icon {{\n",
        "            font-size: 60px;\n",
        "            margin-bottom: 20px;\n",
        "        }}\n",
        "\n",
        "        /* Tooltip styling */\n",
        "        .tooltip {{\n",
        "            position: relative;\n",
        "            display: inline-block;\n",
        "            cursor: help;\n",
        "            color: #3498db;\n",
        "            margin-left: 5px;\n",
        "        }}\n",
        "\n",
        "        .tooltip:hover::after {{\n",
        "            content: attr(data-tooltip);\n",
        "            position: absolute;\n",
        "            top: -30px;\n",
        "            left: 50%;\n",
        "            transform: translateX(-50%);\n",
        "            background: #2c3e50;\n",
        "            color: white;\n",
        "            padding: 5px 10px;\n",
        "            border-radius: 4px;\n",
        "            white-space: nowrap;\n",
        "            font-size: 12px;\n",
        "            z-index: 1000;\n",
        "        }}\n",
        "    </style>\n",
        "</head>\n",
        "<body>\n",
        "    <h1>Chunk Viewer</h1>\n",
        "    <p>Displaying {num_chunks} chunks from Ford car manual</p>\n",
        "\"\"\" # End of initial HTML string\n",
        "\n",
        "    for i, chunk in enumerate(chunks):\n",
        "        chunk_id = chunk.get('id', f'chunk_{i+1}')  # Use f-string for fallback ID\n",
        "        prev_id = chunk.get('prev_chunk_id')\n",
        "        next_id = chunk.get('next_chunk_id')\n",
        "\n",
        "        # Use f-strings for building the rest of the HTML\n",
        "        html_content += f'<div class=\"chunk\" id=\"{chunk_id}\">\\n'\n",
        "        # Header with ID and Nav\n",
        "        html_content += '  <div class=\"chunk-header\">\\n'\n",
        "        html_content += f'    <span class=\"chunk-id\">{chunk_id}</span>\\n'\n",
        "        html_content += '    <span class=\"chunk-nav\">'\n",
        "        if prev_id:\n",
        "            html_content += f'<a href=\"#{prev_id}\">Previous</a>'\n",
        "        if next_id:\n",
        "            html_content += f'<a href=\"#{next_id}\">Next</a>'\n",
        "        html_content += '    </span>\\n'\n",
        "        html_content += '  </div>\\n'\n",
        "\n",
        "        # Metadata\n",
        "        html_content += '  <div class=\"chunk-meta\">\\n'\n",
        "        pages = chunk.get('page_numbers', [])\n",
        "        page_str = f\"{min(pages)}-{max(pages)}\" if pages else \"N/A\"\n",
        "        # Use html.escape for potentially unsafe values if needed, though numbers/context are likely safe\n",
        "        html_content += f\"    <strong>Pages:</strong> {page_str} ({len(pages)}) | \"\n",
        "        html_content += f\"<strong>Length:</strong> {len(chunk.get('text', ''))} chars | \"\n",
        "        html_content += f\"<strong>Context:</strong> {html.escape(chunk.get('breadcrumb_trail', 'N/A'))}<br/>\\n\" # Escape context just in case\n",
        "        if chunk.get('content_type'):\n",
        "            html_content += f\"    <strong>Content Types:</strong> {', '.join(chunk['content_type'])}\\n\"\n",
        "        html_content += '  </div>\\n'\n",
        "\n",
        "        # Features (using f-strings and html.escape where appropriate)\n",
        "        features_html = \"\"\n",
        "        meta = chunk.get('metadata', {})\n",
        "        if meta.get('has_safety'):\n",
        "            features_html += \"<li>Safety Warnings Present</li>\"\n",
        "        if meta.get('systems'):\n",
        "            features_html += f\"<li><strong>Systems:</strong> {', '.join(html.escape(s) for s in meta['systems'])}</li>\"\n",
        "        if meta.get('parts'):\n",
        "            features_html += f\"<li><strong>Parts:</strong> {', '.join(html.escape(p) for p in meta['parts'])}</li>\"\n",
        "            features_html += f\"<li><em>{html.escape(cat)}:</em> {', '.join(html.escape(t) for t in terms)}</li>\"\n",
        "            features_html += \"</ul></li>\"\n",
        "        if chunk.get('figure_references'): # Assuming this comes from chunk, not meta\n",
        "             features_html += f\"<li><strong>Figures:</strong> {', '.join(html.escape(fig) for fig in chunk['figure_references'])}</li>\"\n",
        "        if chunk.get('diagnostic_codes'): # Assuming this comes from chunk, not meta\n",
        "             features_html += f\"<li><strong>DTCs:</strong> {', '.join(html.escape(dtc) for dtc in chunk['diagnostic_codes'])}</li>\"\n",
        "\n",
        "        if chunk.get('related_chunks'):\n",
        "            links = [f'<a href=\"#{rid}\">{rid}</a>' for rid in chunk['related_chunks']]\n",
        "            features_html += f'<li><strong>Related:</strong> {\" \".join(links)}</li>'\n",
        "\n",
        "        if features_html:\n",
        "            html_content += f'<div class=\"features\"><strong>Extracted Features:</strong><ul>{features_html}</ul></div>\\n'\n",
        "\n",
        "        # Text Content - IMPORTANT: Escape the chunk text!\n",
        "        text_display = html.escape(chunk.get('text', ''))\n",
        "        # Highlight overlap prefix after escaping\n",
        "        if chunk.get('has_overlap_prefix'):\n",
        "            import re\n",
        "            # Use non-capturing group for the prefix to simplify replacement\n",
        "            # Match the literal brackets, not regex special characters\n",
        "            prefix_pattern = r'^(\\[Overlap Start\\].*?\\[Overlap End\\]\\\\n\\\\n)'\n",
        "            match = re.match(prefix_pattern, text_display, re.DOTALL)\n",
        "            if match:\n",
        "                prefix = match.group(1)\n",
        "                remaining_text = text_display[len(prefix):]\n",
        "                text_display = f'<span class=\"overlap-prefix\">{prefix}</span>{remaining_text}'\n",
        "\n",
        "        # Basic highlighting for warnings\n",
        "        text_display = text_display.replace(\"⚠️\", '<span class=\"warning\">⚠️</span>')\n",
        "        html_content += f'  <div class=\"chunk-text\"><pre>{text_display}</pre></div>\\n'  # Use <pre> for spacing\n",
        "\n",
        "        html_content += '</div>\\n\\n'\n",
        "\n",
        "    html_content += \"\"\"\n",
        "</body>\n",
        "</html>\n",
        "\"\"\"\n",
        "    try:\n",
        "        with open(filename, 'w', encoding='utf-8') as f:\n",
        "            f.write(html_content)\n",
        "        print(f\"HTML viewer saved to '{filename}'\")\n",
        "        return filename\n",
        "    except Exception as e:\n",
        "        print(f\"Error saving HTML viewer: {e}\")\n",
        "        return None"
      ],
      "metadata": {
        "id": "UBQC-KL5-_8_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "MAIN\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "TCKHEiz9_FPW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import nltk\n",
        "nltk.download('punkt_tab')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Il5zlw9BHdb5",
        "outputId": "9851da1d-e004-4710-c641-c2ea5d74260b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "[nltk_data] Downloading package punkt_tab to /root/nltk_data...\n",
            "[nltk_data]   Unzipping tokenizers/punkt_tab.zip.\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# --- Main Execution ---\n",
        "if 'combined_data_structured' in locals() and combined_data_structured:\n",
        "    print(\"Starting enhanced chunk creation process...\")\n",
        "    # Using process_elements_in_batches even for smaller data for consistency\n",
        "    enriched_final_chunks = process_elements_in_batches(\n",
        "        combined_data=combined_data_structured,\n",
        "        max_chunk_size=MAX_CHUNK_SIZE,\n",
        "        min_chunk_size=MIN_CHUNK_SIZE,\n",
        "        overlap_sentences=OVERLAP_SENTENCE_COUNT\n",
        "        # batch_size can be adjusted based on memory constraints\n",
        "    )\n",
        "\n",
        "    # Perform validation\n",
        "    validate_chunking(enriched_final_chunks)\n",
        "\n",
        "    # Save the final chunks for inspection\n",
        "    output_filename_json = 'final_chunks_output.json'\n",
        "    try:\n",
        "        # Use a custom encoder if you have non-serializable types (like datetime), else standard dump\n",
        "        class CustomEncoder(json.JSONEncoder):\n",
        "            def default(self, obj):\n",
        "                if isinstance(obj, (datetime.datetime, datetime.date)):\n",
        "                    return obj.isoformat()\n",
        "                # Add other type handlers if needed\n",
        "                return json.JSONEncoder.default(self, obj)\n",
        "\n",
        "        with open(output_filename_json, 'w', encoding='utf-8') as f:\n",
        "            json.dump(enriched_final_chunks, f, indent=2, ensure_ascii=False, cls=CustomEncoder)\n",
        "        print(f\"\\nFinal chunks saved to '{output_filename_json}'\")\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"Error saving final chunks to JSON: {e}\")\n",
        "\n",
        "    # Generate HTML viewer\n",
        "    create_html_viewer(enriched_final_chunks, \"enriched_chunks_viewer.html\")\n",
        "\n",
        "else:\n",
        "    print(\"Error: 'combined_data_structured' not found or is empty. Please ensure it's loaded correctly before running this script.\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BeDlbGnr_HTr",
        "outputId": "23a333bf-2e5b-4766-c3c6-6717447c906f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Starting enhanced chunk creation process...\n",
            "Processing 8172 elements in batches of 1000...\n",
            "Manual spans 486 pages across 8172 elements.\n",
            "\n",
            "Processing batch 1/9: Elements 1-1000 (Pages 1-70)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 2790.26element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 99 chunks. Total chunks so far: 99\n",
            "\n",
            "Processing batch 2/9: Elements 1001-2000 (Pages 70-128)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 3530.58element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 85 chunks. Total chunks so far: 184\n",
            "\n",
            "Processing batch 3/9: Elements 2001-3000 (Pages 128-184)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 3998.52element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 84 chunks. Total chunks so far: 268\n",
            "\n",
            "Processing batch 4/9: Elements 3001-4000 (Pages 184-244)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 3109.36element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 96 chunks. Total chunks so far: 364\n",
            "\n",
            "Processing batch 5/9: Elements 4001-5000 (Pages 244-304)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 3058.57element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 98 chunks. Total chunks so far: 462\n",
            "\n",
            "Processing batch 6/9: Elements 5001-6000 (Pages 304-346)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 6943.07element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 49 chunks. Total chunks so far: 511\n",
            "\n",
            "Processing batch 7/9: Elements 6001-7000 (Pages 346-422)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 2587.51element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 114 chunks. Total chunks so far: 625\n",
            "\n",
            "Processing batch 8/9: Elements 7001-8000 (Pages 422-475)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 1000/1000 [00:00<00:00, 4323.79element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 71 chunks. Total chunks so far: 696\n",
            "\n",
            "Processing batch 9/9: Elements 8001-8172 (Pages 475-486)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Creating chunks: 100%|██████████| 172/172 [00:00<00:00, 2706.25element/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Batch complete. Generated 18 chunks. Total chunks so far: 714\n",
            "\n",
            "Finalizing all chunks...\n",
            "Processing complete. Total chunks: 714\n",
            "\n",
            "--- Starting Chunk Validation ---\n",
            "Total Chunks Generated: 714\n",
            "\n",
            "Chunk Size Analysis:\n",
            "  Average Size: 1240 characters\n",
            "  Min Size: 300 | Max Size: 2119\n",
            "  Chunks below Min (500): 43 (6.0%)\n",
            "  Chunks above Max (1800): 19 (2.7%)\n",
            "  Chunks significantly above Max (2160): 0 (0.0%)\n",
            "  WARNING: Still found 43 chunks below 500 chars. Review merge logic or MIN_CHUNK_SIZE.\n",
            "\n",
            "Overlap Analysis (Expected sentences: 1):\n",
            "  Chunks with Overlap Prefix: 0 (0.0%)\n",
            "\n",
            "Context Coherence Check (Sample):\n",
            "  Context appears generally coherent in sample based on hierarchy.\n",
            "\n",
            "Readability Check:\n",
            "  Please manually review chunks in the generated HTML file ('enriched_chunks_viewer.html'),\n",
            "  especially those containing lists, procedures, tables, or warnings,\n",
            "  to ensure text flows correctly and formatting is acceptable.\n",
            "\n",
            "--- Validation Complete ---\n",
            "\n",
            "Final chunks saved to 'final_chunks_output.json'\n",
            "HTML viewer saved to 'enriched_chunks_viewer.html'\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 4. Create Embeddings with Vertex AI and Store in MongoDB Atlas Cluster"
      ],
      "metadata": {
        "id": "QMEuRiTI2lmI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now send our final chunks to a **Vertex AI embedding model** and store them with their respective embeddings in MongoDB in the following steps:\n",
        "\n",
        "1. Generate embeddings using Vertex AI\n",
        "2. Format the data for MongoDB\n",
        "3. Store the chunks and embeddings in MongoDB"
      ],
      "metadata": {
        "id": "v98BUwzx23aY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Install necessary packages\n",
        "!pip install pymongo\n",
        "!pip install pymongo google-cloud-aiplatform"
      ],
      "metadata": {
        "id": "D54pf_rGbfno"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "import os\n",
        "import datetime\n",
        "import time\n",
        "from google.cloud import aiplatform\n",
        "from typing import List, Dict, Any, Optional\n",
        "\n",
        "def safe_mongodb_imports():\n",
        "    \"\"\"Import MongoDB-related modules safely, avoiding the SON conflict.\"\"\"\n",
        "    import json\n",
        "    import pymongo # Import the main package\n",
        "    # Create local references to needed functions and classes from pymongo\n",
        "    MongoClient = pymongo.MongoClient\n",
        "    # Import json_util from bson (which should now be pymongo's bson)\n",
        "    from bson import json_util\n",
        "\n",
        "    # Return the necessary components\n",
        "    return MongoClient, json_util\n",
        "\n",
        "# --- Use the safe import function defined previously ---\n",
        "try:\n",
        "    MongoClient, json_util = safe_mongodb_imports()\n",
        "    print(\"MongoDB modules imported safely.\")\n",
        "except NameError:\n",
        "    print(\"Error: 'safe_mongodb_imports' not defined. Please run the cell that defines it (e.g., Bdbbnyonzsma).\")\n",
        "    # Fallback to direct imports, hoping the environment is fixed\n",
        "    from pymongo import MongoClient, ASCENDING, IndexModel\n",
        "    from pymongo.errors import ConnectionFailure, OperationFailure\n",
        "    from bson import json_util\n",
        "\n",
        "\n",
        "# --- Configuration ---\n",
        "project_id = \"ENTER_PROJECT_ID\" # Your GCP Project ID\n",
        "location = \"ENTER_LOCATION\" # GCP location for Vertex AI\n",
        "model_name = \"text-embedding-005\" # Or the specific model you use\n",
        "# Embedding dimension for the model (e.g., text-embedding-005 uses 768)\n",
        "EXPECTED_EMBEDDING_DIMENSION = 768\n",
        "\n",
        "mongodb_uri = \"YOUR_URI\" #enter correct uri here\n",
        "db_name = \"DB_NAME\" # Your DB name\n",
        "collection_name = \"COLLECTION_NAME\" # Your Collection name"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3LeInXqoattw",
        "outputId": "071f3d18-0d3d-414b-c717-bb8adc3768e1"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "MongoDB modules imported safely.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "GENERATE EMBEDDINGS\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "WGBUoLo-cIZF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Generate embeddings using Vertex AI with batching and error handling\n",
        "def generate_embeddings_vertex(\n",
        "    chunks: List[Dict[str, Any]],\n",
        "    project_id: str,\n",
        "    location: str = \"us-central1\",\n",
        "    model_id: str = \"text-embedding-005\",\n",
        "    task_type: str = \"RETRIEVAL_DOCUMENT\",\n",
        "    output_dim: Optional[int] = None,\n",
        "    batch_size: int = 5\n",
        "):\n",
        "    \"\"\"\n",
        "    Generate embeddings using the Vertex AI Text Embedding API.\n",
        "\n",
        "    Args:\n",
        "        chunks: List of chunks to generate embeddings for\n",
        "        project_id: GCP project ID\n",
        "        location: GCP location (e.g., \"us-central1\")\n",
        "        model_id: Model ID (e.g., \"text-embedding-005\")\n",
        "        task_type: Task type for embedding (e.g., \"RETRIEVAL_DOCUMENT\")\n",
        "        output_dim: Optional output dimensionality (between 1 and max dimensions)\n",
        "        batch_size: Number of chunks to process in each batch\n",
        "\n",
        "    Returns:\n",
        "        List of embeddings corresponding to the chunks\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Import the Vertex AI SDK modules\n",
        "        from vertexai.language_models import TextEmbeddingModel, TextEmbeddingInput\n",
        "    except ImportError:\n",
        "        raise ImportError(\n",
        "            \"The vertexai package is required. Install with: pip install google-cloud-aiplatform\"\n",
        "        )\n",
        "\n",
        "    # Initialize Vertex AI\n",
        "    import vertexai\n",
        "    vertexai.init(project=project_id, location=location)\n",
        "\n",
        "    # Load the text embedding model\n",
        "    model = TextEmbeddingModel.from_pretrained(model_id)\n",
        "    print(f\"Successfully loaded embedding model: {model_id}\")\n",
        "\n",
        "    all_embeddings = []\n",
        "\n",
        "    # Process chunks in batches\n",
        "    for i in range(0, len(chunks), batch_size):\n",
        "        batch = chunks[i:i+batch_size]\n",
        "        print(f\"Processing embedding batch {i//batch_size + 1}/{(len(chunks)-1)//batch_size + 1}\")\n",
        "\n",
        "        # Extract text from each chunk\n",
        "        batch_texts = [chunk['text'] if isinstance(chunk, dict) else chunk for chunk in batch]\n",
        "\n",
        "        # Create TextEmbeddingInput objects\n",
        "        inputs = [TextEmbeddingInput(text=text, task_type=task_type) for text in batch_texts]\n",
        "\n",
        "        try:\n",
        "            # Set up the keyword arguments for dimensionality\n",
        "            kwargs = {}\n",
        "            if output_dim is not None:\n",
        "                kwargs['output_dimensionality'] = output_dim\n",
        "\n",
        "            # Get embeddings for the batch\n",
        "            embeddings = model.get_embeddings(inputs, **kwargs)\n",
        "\n",
        "            # Extract the embedding vectors\n",
        "            batch_embeddings = [emb.values for emb in embeddings]\n",
        "            all_embeddings.extend(batch_embeddings)\n",
        "\n",
        "            # Add short delay to avoid rate limiting\n",
        "            if i + batch_size < len(chunks):\n",
        "                time.sleep(0.1)\n",
        "\n",
        "        except Exception as e:\n",
        "            print(f\"Error generating embeddings for batch {i//batch_size + 1}: {str(e)}\")\n",
        "            # Add None for failed embeddings\n",
        "            all_embeddings.extend([None] * len(batch))\n",
        "\n",
        "    # Check for any failed embeddings\n",
        "    failed_count = sum(1 for emb in all_embeddings if emb is None)\n",
        "    if failed_count > 0:\n",
        "        print(f\"Warning: {failed_count} chunks failed to generate embeddings\")\n",
        "\n",
        "    return all_embeddings"
      ],
      "metadata": {
        "id": "LV36acso2xR_"
      },
      "outputs": [],
      "execution_count": null
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "---\n",
        "\n",
        "CREATE MONGODB DOCUMENTS\n",
        "\n",
        "\n",
        "\n",
        "---\n",
        "\n"
      ],
      "metadata": {
        "id": "Gqk-M6TScMZR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create MongoDB documents with chunks and embeddings\n",
        "def create_mongodb_documents(chunks, embeddings):\n",
        "    \"\"\"\n",
        "    Create MongoDB documents from chunks and embeddings.\n",
        "\n",
        "    Args:\n",
        "        chunks: List of chunk documents\n",
        "        embeddings: List of embeddings for the chunks\n",
        "\n",
        "    Returns:\n",
        "        List of documents ready for MongoDB insertion\n",
        "    \"\"\"\n",
        "    documents = []\n",
        "    for chunk, embedding in zip(chunks, embeddings):\n",
        "        # Skip if embedding generation failed\n",
        "        if embedding is None:\n",
        "            continue\n",
        "\n",
        "        # Create document with all chunk fields and embedding\n",
        "        document = chunk.copy() if isinstance(chunk, dict) else {\"text\": chunk}\n",
        "        document[\"embedding\"] = embedding\n",
        "        document[\"embedding_timestamp\"] = datetime.datetime.now().isoformat()\n",
        "\n",
        "        # Sanitize document for MongoDB\n",
        "        _sanitize_document(document)\n",
        "\n",
        "        documents.append(document)\n",
        "\n",
        "    return documents\n",
        "\n",
        "# Helper function to sanitize documents for MongoDB compatibility\n",
        "def _sanitize_document(doc):\n",
        "    \"\"\"Clean document for MongoDB compatibility\"\"\"\n",
        "    # Replace '.' in keys with '_' (MongoDB doesn't allow dots in field names)\n",
        "    keys_to_update = [k for k in doc.keys() if '.' in k]\n",
        "    for key in keys_to_update:\n",
        "        new_key = key.replace('.', '_')\n",
        "        doc[new_key] = doc.pop(key)\n",
        "\n",
        "    # Handle nested dictionaries\n",
        "    for key, value in doc.items():\n",
        "        if isinstance(value, dict):\n",
        "            _sanitize_document(value)\n",
        "        elif isinstance(value, list):\n",
        "            for item in value:\n",
        "                if isinstance(item, dict):\n",
        "                    _sanitize_document(item)\n",
        "\n",
        "# Store documents in MongoDB with proper connection handling\n",
        "def store_in_mongodb(documents, mongodb_uri, db_name, collection_name):\n",
        "    \"\"\"\n",
        "    Store documents in MongoDB with error handling and connection management.\n",
        "\n",
        "    Args:\n",
        "        documents: List of documents to store\n",
        "        mongodb_uri: MongoDB connection URI\n",
        "        db_name: Database name\n",
        "        collection_name: Collection name\n",
        "\n",
        "    Returns:\n",
        "        Result of the insert operation\n",
        "    \"\"\"\n",
        "    if not documents:\n",
        "        print(\"No documents to store in MongoDB\")\n",
        "        return None\n",
        "\n",
        "    # Validate MongoDB connection parameters\n",
        "    if not mongodb_uri or not db_name or not collection_name:\n",
        "        raise ValueError(\"MongoDB URI, database name, and collection name are required\")\n",
        "\n",
        "    client = None\n",
        "    try:\n",
        "        # Connect to MongoDB with a timeout\n",
        "        client = MongoClient(mongodb_uri, serverSelectionTimeoutMS=5000)\n",
        "\n",
        "        # Test the connection\n",
        "        client.admin.command('ping')\n",
        "\n",
        "        db = client[db_name]\n",
        "        collection = db[collection_name]\n",
        "\n",
        "        # Insert documents with a reasonable batch size\n",
        "        batch_size = min(1000, len(documents))  # MongoDB has a 16MB limit per batch\n",
        "        result = None\n",
        "\n",
        "        for i in range(0, len(documents), batch_size):\n",
        "            batch = documents[i:i+batch_size]\n",
        "            batch_result = collection.insert_many(batch)\n",
        "\n",
        "            if result is None:\n",
        "                result = batch_result\n",
        "            else:\n",
        "                # Combine inserted IDs\n",
        "                result.inserted_ids.extend(batch_result.inserted_ids)\n",
        "\n",
        "            print(f\"Inserted batch {i//batch_size + 1}/{(len(documents)-1)//batch_size + 1} ({len(batch)} documents)\")\n",
        "\n",
        "        return result\n",
        "    except ConnectionFailure as e:\n",
        "        print(f\"MongoDB connection error: {str(e)}\")\n",
        "        raise\n",
        "    except Exception as e:\n",
        "        print(f\"Error storing documents in MongoDB: {str(e)}\")\n",
        "        raise\n",
        "    finally:\n",
        "        # Ensure the client connection is closed\n",
        "        if client:\n",
        "            client.close()\n",
        "            print(\"MongoDB connection closed\")\n",
        "\n",
        "# Main pipeline with comprehensive error handling\n",
        "def process_chunks_to_mongodb(\n",
        "    final_chunks,\n",
        "    project_id,\n",
        "    location,\n",
        "    model_id,\n",
        "    mongodb_uri,\n",
        "    db_name,\n",
        "    collection_name,\n",
        "    task_type=\"RETRIEVAL_DOCUMENT\",\n",
        "    output_dim=None\n",
        "):\n",
        "    \"\"\"\n",
        "    Process chunks, generate embeddings, and store in MongoDB.\n",
        "\n",
        "    Args:\n",
        "        final_chunks: List of processed chunks\n",
        "        project_id: GCP project ID\n",
        "        location: GCP location\n",
        "        model_id: Model ID for embeddings\n",
        "        mongodb_uri: MongoDB connection URI\n",
        "        db_name: MongoDB database name\n",
        "        collection_name: MongoDB collection name\n",
        "        task_type: Task type for embeddings\n",
        "        output_dim: Optional output dimensionality\n",
        "\n",
        "    Returns:\n",
        "        Result of the MongoDB insert operation\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Generate embeddings\n",
        "        print(f\"Generating embeddings for {len(final_chunks)} chunks...\")\n",
        "        embeddings = generate_embeddings_vertex(\n",
        "            chunks=final_chunks,\n",
        "            project_id=project_id,\n",
        "            location=location,\n",
        "            model_id=model_id,\n",
        "            task_type=task_type,\n",
        "            output_dim=output_dim\n",
        "        )\n",
        "        print(f\"Generated {len(embeddings)} embeddings\")\n",
        "\n",
        "        # Create MongoDB documents\n",
        "        print(\"Creating MongoDB documents...\")\n",
        "        documents = create_mongodb_documents(final_chunks, embeddings)\n",
        "        print(f\"Created {len(documents)} documents\")\n",
        "\n",
        "        # Store in MongoDB\n",
        "        print(f\"Storing documents in MongoDB ({db_name}.{collection_name})...\")\n",
        "        result = store_in_mongodb(documents, mongodb_uri, db_name, collection_name)\n",
        "\n",
        "        return result\n",
        "    except Exception as e:\n",
        "        print(f\"Error in processing pipeline: {str(e)}\")\n",
        "        raise"
      ],
      "metadata": {
        "id": "MExrToHerKfI"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}